[Logprob] Serve input-logprob temporaries from CUDA-graph-pool dead space (#40038)
Co-authored-by: cctry <csycfl@gmail.com>
This commit is contained in:
metamergebot
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
"""Logprob memory budgets must preserve TP collectives and graph replay safety."""
|
||||
|
||||
import sys
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.logprob_processor import InputLogprobProcessor
|
||||
from sglang.srt.model_executor.runner_utils import pool
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=60, stage="base-b", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
def _run_rank(rank, rendezvous):
|
||||
torch.cuda.set_device(rank)
|
||||
dist.init_process_group(
|
||||
"gloo",
|
||||
init_method=f"file://{rendezvous}",
|
||||
rank=rank,
|
||||
world_size=2,
|
||||
timeout=timedelta(seconds=60),
|
||||
)
|
||||
nccl_group = dist.new_group(backend="nccl", timeout=timedelta(seconds=60))
|
||||
try:
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda", rank)
|
||||
rows, vocab = 2048, 202752
|
||||
states = torch.randint(-2, 3, (rows, 64), device=device).to(torch.bfloat16)
|
||||
weight = torch.randint(-2, 3, (vocab, 64), device=device).to(torch.bfloat16)
|
||||
local_weight = weight.chunk(2)[rank].contiguous()
|
||||
token_ids = torch.randint(vocab, (rows,), device=device)
|
||||
sample_rows = [rows // 2 - 1, rows - 1]
|
||||
metadata = SimpleNamespace(
|
||||
sample_indices_cpu=sample_rows,
|
||||
input_logprob_indices_cpu=list(range(rows)),
|
||||
extend_return_top_logprob=False,
|
||||
extend_token_ids_logprob=False,
|
||||
top_logprobs_nums=None,
|
||||
extend_logprob_pruned_lens_cpu=[rows // 2, rows // 2],
|
||||
extend_input_logprob_token_ids_gpu=token_ids,
|
||||
token_ids_logprobs=None,
|
||||
)
|
||||
chunk_rows = []
|
||||
allocations_borrowed = []
|
||||
runs = []
|
||||
|
||||
def get_logits(chunk, *_args, **_kwargs):
|
||||
chunk_rows.append(chunk.shape[0])
|
||||
local = torch.mm(chunk, local_weight.T)
|
||||
gathered = torch.empty(
|
||||
(2 * chunk.shape[0], vocab // 2), device=device, dtype=local.dtype
|
||||
)
|
||||
dist.all_gather_into_tensor(gathered, local, group=nccl_group)
|
||||
reshaped = (
|
||||
gathered.reshape(2, chunk.shape[0], vocab // 2)
|
||||
.movedim(0, 1)
|
||||
.reshape(chunk.shape[0], vocab)
|
||||
)
|
||||
converted = reshaped.float()
|
||||
allocations_borrowed.extend(
|
||||
any(
|
||||
lo <= tensor.data_ptr()
|
||||
and tensor.data_ptr() + tensor.nbytes <= lo + size
|
||||
for lo, size in runs
|
||||
)
|
||||
for tensor in (local, gathered, reshaped, converted)
|
||||
)
|
||||
return converted
|
||||
|
||||
processor = InputLogprobProcessor(vocab, chunking_group=dist.group.WORLD)
|
||||
args = dict(
|
||||
pruned_states=states,
|
||||
sample_indices=torch.tensor(sample_rows, device=device),
|
||||
input_logprob_indices=torch.arange(rows, device=device),
|
||||
token_to_seq_idx=[0] * (rows // 2) + [1] * (rows // 2),
|
||||
lm_head=None,
|
||||
get_logits_fn=get_logits,
|
||||
logits_metadata=metadata,
|
||||
)
|
||||
for fast in (False, True):
|
||||
processor.enable_fast_input_logprobs = fast
|
||||
processor.enable_logprobs_chunk = False
|
||||
reference, reference_sampled = processor.forward(**args)
|
||||
reference_logprobs = reference.token_logprobs.cpu()
|
||||
processor.enable_logprobs_chunk = True
|
||||
for chunk_size, disabled_rank, borrowing in (
|
||||
(128, None, True),
|
||||
(256, None, False),
|
||||
(128, 1, False),
|
||||
):
|
||||
processor.logprobs_chunk_size = chunk_size
|
||||
state = pool.GraphPoolBorrowState()
|
||||
handle = torch.cuda.graph_pool_handle()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
seed = torch.zeros(8, device=device)
|
||||
stream = torch.cuda.Stream()
|
||||
with (
|
||||
torch.cuda.stream(stream),
|
||||
torch.cuda.graph(graph, pool=handle, stream=stream),
|
||||
):
|
||||
transient = torch.empty(
|
||||
(512 + rank * 256) << 20, dtype=torch.uint8, device=device
|
||||
)
|
||||
transient.fill_(7)
|
||||
keep = seed + 1
|
||||
del transient
|
||||
torch.cuda.synchronize()
|
||||
state.disabled_reason = (
|
||||
"test fallback" if rank == disabled_rank else None
|
||||
)
|
||||
with (
|
||||
pool.get_resources().override(graph_pool_borrow=state),
|
||||
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
|
||||
patch.object(
|
||||
pool, "get_global_graph_memory_pool", return_value=handle
|
||||
),
|
||||
patch.object(
|
||||
torch.cuda,
|
||||
"mem_get_info",
|
||||
side_effect=AssertionError(
|
||||
"Chunk sizing must not query heap headroom"
|
||||
),
|
||||
),
|
||||
):
|
||||
chunk_rows.clear()
|
||||
allocations_borrowed.clear()
|
||||
runs = pool.find_free_graph_pool_runs(handle)
|
||||
result, sampled = processor.forward(**args)
|
||||
with pool.graph_pool_replay_scope():
|
||||
graph.replay()
|
||||
if result.input_copy_done is not None:
|
||||
result.input_copy_done.synchronize()
|
||||
assert torch.equal(result.token_logprobs.cpu(), reference_logprobs)
|
||||
assert torch.equal(sampled, reference_sampled)
|
||||
assert (result.input_copy_done is not None) == borrowing
|
||||
assert all(
|
||||
borrowed == borrowing for borrowed in allocations_borrowed
|
||||
)
|
||||
all_chunk_rows = [None, None]
|
||||
dist.all_gather_object(all_chunk_rows, chunk_rows)
|
||||
assert all_chunk_rows[0] == all_chunk_rows[1]
|
||||
assert chunk_rows == [chunk_size] * (rows // chunk_size)
|
||||
pool._teardown_borrow_pool()
|
||||
del graph, keep
|
||||
finally:
|
||||
dist.destroy_process_group(nccl_group)
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices")
|
||||
def test_logprob_chunks_share_tp_budget_and_survive_replay(tmp_path):
|
||||
mp.spawn(_run_rank, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -8,9 +8,11 @@ the scheduler asserts on.
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.logprob_processor import InputLogprobProcessor
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.logprob_test_utils import coverage_cases
|
||||
@@ -52,6 +54,8 @@ def _build_batch(seq_specs, with_token_ids):
|
||||
lp_pt += rows
|
||||
pruned_lens.append(n_lp)
|
||||
metadata = SimpleNamespace(
|
||||
sample_indices_cpu=sample_indices,
|
||||
input_logprob_indices_cpu=input_logprob_indices,
|
||||
extend_return_top_logprob=True,
|
||||
extend_token_ids_logprob=with_token_ids,
|
||||
top_logprobs_nums=[TOPK_CYCLE[i % 3] for i in range(len(seq_specs))],
|
||||
@@ -136,6 +140,24 @@ class TestLogprobChunkStitching(CustomTestCase):
|
||||
def test_token_ids_logprobs_stitching(self):
|
||||
self._sweep(with_token_ids=True)
|
||||
|
||||
def test_finalizing_input_logprobs_preserves_request_boundaries(self):
|
||||
rows = [torch.tensor([[1.0], [2.0]]), torch.tensor([[3.0]])]
|
||||
copy_done = Mock()
|
||||
output = LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
input_token_ids_logprobs_val=[[rows[0][:1], rows[0][1:]], [rows[1]]],
|
||||
input_logprobs_copy_done=copy_done,
|
||||
)
|
||||
output.finalize_input_logprobs()
|
||||
self.assertEqual(output.input_token_ids_logprobs_val, [[[1.0], [2.0]], [[3.0]]])
|
||||
copy_done.synchronize.assert_called_once_with()
|
||||
self.assertIsNone(output.input_logprobs_copy_done)
|
||||
|
||||
# Multi-item scoring returns one tensor per request, with no borrow event.
|
||||
output.input_token_ids_logprobs_val = rows
|
||||
output.finalize_input_logprobs()
|
||||
self.assertIs(output.input_token_ids_logprobs_val, rows)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -56,6 +56,8 @@ def _build_batch(seq_specs, dtype, vocab=VOCAB):
|
||||
lp_pt += rows
|
||||
pruned_lens.append(n_lp)
|
||||
metadata = SimpleNamespace(
|
||||
sample_indices_cpu=sample_indices,
|
||||
input_logprob_indices_cpu=input_logprob_indices,
|
||||
extend_return_top_logprob=True,
|
||||
extend_token_ids_logprob=True,
|
||||
top_logprobs_nums=[TOPK_CYCLE[i % 3] for i in range(len(seq_specs))],
|
||||
|
||||
@@ -9,6 +9,11 @@ from unittest.mock import MagicMock, Mock, patch
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.logprob_processor import InputLogprobProcessor
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
)
|
||||
from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
|
||||
FullCudaGraphBackend,
|
||||
)
|
||||
@@ -339,6 +344,242 @@ class TestGraphPoolBorrow(CustomTestCase):
|
||||
self.assertEqual(torch.cuda.memory_reserved(device_id), reserved_before)
|
||||
del graph, y
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
|
||||
def test_input_logprobs_survive_replay_with_growing_chunks(self):
|
||||
handle = torch.cuda.graph_pool_handle()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
seed = torch.zeros(8, device="cuda")
|
||||
stream = torch.cuda.Stream()
|
||||
with (
|
||||
torch.cuda.stream(stream),
|
||||
torch.cuda.graph(graph, pool=handle, stream=stream),
|
||||
):
|
||||
transient = torch.empty(200 << 20, dtype=torch.uint8, device="cuda")
|
||||
transient.fill_(7)
|
||||
keep = seed + 1
|
||||
del transient
|
||||
torch.cuda.synchronize()
|
||||
|
||||
processor = InputLogprobProcessor(vocab_size=4096)
|
||||
processor.enable_fast_input_logprobs = False
|
||||
processor.enable_logprobs_chunk = True
|
||||
with (
|
||||
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
|
||||
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
|
||||
):
|
||||
# Chunk sizes grow without a teardown in between, so each borrow
|
||||
# must still fit after the previous iterations' carves.
|
||||
for rows, chunk_size in (
|
||||
(6, 4),
|
||||
(1500, 1500),
|
||||
(2000, 2000),
|
||||
(2500, 2500),
|
||||
(2600, 2600),
|
||||
):
|
||||
with self.subTest(rows=rows, chunk_size=chunk_size):
|
||||
logits = torch.randn(rows, 4096, device="cuda")
|
||||
token_ids = torch.randint(0, 4096, (rows,), device="cuda")
|
||||
split = rows // 2
|
||||
sample_indices = [split - 1, rows - 1]
|
||||
metadata = SimpleNamespace(
|
||||
sample_indices_cpu=sample_indices,
|
||||
input_logprob_indices_cpu=list(range(rows)),
|
||||
extend_return_top_logprob=True,
|
||||
extend_token_ids_logprob=True,
|
||||
top_logprobs_nums=[2, 3],
|
||||
extend_logprob_pruned_lens_cpu=[split, rows - split],
|
||||
extend_input_logprob_token_ids_gpu=token_ids,
|
||||
token_ids_logprobs=[[0, 7], [4]],
|
||||
)
|
||||
processor.logprobs_chunk_size = chunk_size
|
||||
get_logits = Mock(side_effect=lambda states, *_args, **_kw: states)
|
||||
result, sampled = processor.forward(
|
||||
pruned_states=logits,
|
||||
sample_indices=torch.tensor(sample_indices, device="cuda"),
|
||||
input_logprob_indices=torch.arange(rows, device="cuda"),
|
||||
token_to_seq_idx=[0] * split + [1] * (rows - split),
|
||||
lm_head=None,
|
||||
get_logits_fn=get_logits,
|
||||
logits_metadata=metadata,
|
||||
)
|
||||
self.assertIsNotNone(result.input_copy_done)
|
||||
self.assertTrue(result.token_logprobs.is_pinned())
|
||||
self.assertEqual(
|
||||
[call.args[0].shape[0] for call in get_logits.call_args_list],
|
||||
[min(chunk_size, rows - i) for i in range(0, rows, chunk_size)],
|
||||
)
|
||||
output = LogitsProcessorOutput(next_token_logits=sampled)
|
||||
result.write_input_to(output)
|
||||
with pool.graph_pool_replay_scope():
|
||||
graph.replay()
|
||||
SchedulerBatchResultProcessor.move_logprobs_to_cpu(
|
||||
None,
|
||||
batch=SimpleNamespace(return_logprob=True),
|
||||
logits_output=output,
|
||||
)
|
||||
expected = torch.log_softmax(logits, dim=-1)
|
||||
self.assertEqual(
|
||||
output.input_token_logprobs,
|
||||
tuple(
|
||||
expected[
|
||||
torch.arange(rows, device="cuda"), token_ids
|
||||
].tolist()
|
||||
),
|
||||
)
|
||||
self.assertTrue(torch.equal(sampled, logits[sample_indices]))
|
||||
for i, (lo, hi) in enumerate(((0, split), (split, rows))):
|
||||
values, indices = expected[lo:hi].topk(
|
||||
metadata.top_logprobs_nums[i]
|
||||
)
|
||||
self.assertEqual(
|
||||
output.input_top_logprobs_val[i], values.tolist()
|
||||
)
|
||||
self.assertEqual(
|
||||
output.input_top_logprobs_idx[i], indices.tolist()
|
||||
)
|
||||
self.assertEqual(
|
||||
output.input_token_ids_logprobs_val[i],
|
||||
expected[lo:hi, metadata.token_ids_logprobs[i]].tolist(),
|
||||
)
|
||||
self.assertIsNone(output.input_logprobs_copy_done)
|
||||
pool._teardown_borrow_pool()
|
||||
del graph, keep
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
|
||||
def test_borrow_preserves_planned_chunks(self):
|
||||
"""Graph capacity controls borrowing without changing LM-head shapes."""
|
||||
rows, vocab = 8192, 4096
|
||||
states = torch.randint(-2, 3, (rows, 64), device="cuda").to(torch.bfloat16)
|
||||
weight = torch.randint(-2, 3, (vocab, 64), device="cuda").to(torch.bfloat16)
|
||||
logits = torch.mm(states, weight.T).float()
|
||||
expected = torch.log_softmax(logits, dim=-1)
|
||||
token_ids = torch.randint(0, vocab, (rows,), device="cuda")
|
||||
metadata = SimpleNamespace(
|
||||
sample_indices_cpu=[rows - 1],
|
||||
input_logprob_indices_cpu=list(range(rows)),
|
||||
extend_return_top_logprob=False,
|
||||
extend_token_ids_logprob=False,
|
||||
top_logprobs_nums=None,
|
||||
extend_logprob_pruned_lens_cpu=[rows],
|
||||
extend_input_logprob_token_ids_gpu=token_ids,
|
||||
token_ids_logprobs=None,
|
||||
)
|
||||
processor = InputLogprobProcessor(vocab_size=vocab)
|
||||
processor.enable_fast_input_logprobs = False
|
||||
processor.enable_logprobs_chunk = True
|
||||
processor.logprobs_chunk_size = rows
|
||||
|
||||
def run_with_pool_run_of(nbytes, disabled, lm_head=None):
|
||||
handle = torch.cuda.graph_pool_handle()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
seed = torch.zeros(8, device="cuda")
|
||||
stream = torch.cuda.Stream()
|
||||
with (
|
||||
torch.cuda.stream(stream),
|
||||
torch.cuda.graph(graph, pool=handle, stream=stream),
|
||||
):
|
||||
transient = torch.empty(nbytes, dtype=torch.uint8, device="cuda")
|
||||
keep = seed + 1
|
||||
del transient
|
||||
torch.cuda.synchronize()
|
||||
allocations_borrowed = []
|
||||
chunk_rows = []
|
||||
runs = pool.find_free_graph_pool_runs(handle)
|
||||
|
||||
def get_logits(chunk, *_args, **_kwargs):
|
||||
chunk_rows.append(chunk.shape[0])
|
||||
projected = torch.mm(chunk, weight.T)
|
||||
converted = projected.float()
|
||||
allocations_borrowed.extend(
|
||||
any(
|
||||
lo <= tensor.data_ptr()
|
||||
and tensor.data_ptr() + tensor.nbytes <= lo + size
|
||||
for lo, size in runs
|
||||
)
|
||||
for tensor in (projected, converted)
|
||||
)
|
||||
return converted
|
||||
|
||||
with (
|
||||
envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True),
|
||||
patch.object(pool, "get_global_graph_memory_pool", return_value=handle),
|
||||
patch.object(
|
||||
self.state, "disabled_reason", "test fallback" if disabled else None
|
||||
),
|
||||
patch.object(
|
||||
torch.cuda,
|
||||
"mem_get_info",
|
||||
side_effect=AssertionError(
|
||||
"Chunk sizing must not query heap headroom"
|
||||
),
|
||||
),
|
||||
):
|
||||
# Precarve outside the measurement so the peak covers the forward.
|
||||
with pool.borrow_graph_pool(user="warmup"):
|
||||
pass
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
base = torch.cuda.memory_stats()["allocated_bytes.all.current"]
|
||||
result, sampled = processor.forward(
|
||||
pruned_states=states,
|
||||
sample_indices=torch.tensor([rows - 1], device="cuda"),
|
||||
input_logprob_indices=torch.arange(rows, device="cuda"),
|
||||
token_to_seq_idx=[0] * rows,
|
||||
lm_head=lm_head,
|
||||
get_logits_fn=get_logits,
|
||||
logits_metadata=metadata,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] - base
|
||||
with pool.graph_pool_replay_scope():
|
||||
graph.replay()
|
||||
pool._teardown_borrow_pool()
|
||||
del graph, keep
|
||||
return result, sampled, peak, allocations_borrowed, chunk_rows
|
||||
|
||||
expected_token_logprobs = expected[
|
||||
torch.arange(rows, device="cuda"), token_ids
|
||||
].cpu()
|
||||
|
||||
for chunk_size, nbytes, disabled, borrowing in (
|
||||
(rows, 512 << 20, False, True),
|
||||
(rows, 96 << 20, False, False),
|
||||
(rows, 16 << 20, False, False),
|
||||
(rows, 512 << 20, True, False),
|
||||
(32, 66 << 20, False, True),
|
||||
(32, 66 << 20, True, False),
|
||||
):
|
||||
with self.subTest(chunk_size=chunk_size, nbytes=nbytes, disabled=disabled):
|
||||
processor.logprobs_chunk_size = chunk_size
|
||||
result, sampled, peak, borrowed, chunk_rows = run_with_pool_run_of(
|
||||
nbytes, disabled
|
||||
)
|
||||
self.assertEqual(
|
||||
chunk_rows,
|
||||
[min(chunk_size, rows - i) for i in range(0, rows, chunk_size)],
|
||||
)
|
||||
self.assertTrue(all(value == borrowing for value in borrowed))
|
||||
if borrowing:
|
||||
self.assertLessEqual(
|
||||
peak, 2 * max(chunk_rows) * vocab * 4 + (2 << 20)
|
||||
)
|
||||
self.assertIsNotNone(result.input_copy_done)
|
||||
result.input_copy_done.synchronize()
|
||||
self.assertTrue(result.token_logprobs.is_pinned())
|
||||
else:
|
||||
self.assertIsNone(result.input_copy_done)
|
||||
self.assertTrue(result.token_logprobs.is_cuda)
|
||||
self.assertTrue(
|
||||
torch.equal(result.token_logprobs.cpu(), expected_token_logprobs)
|
||||
)
|
||||
self.assertTrue(torch.equal(sampled, logits[-1:]))
|
||||
|
||||
# LoRA has already prepared adapter metadata for each configured pass.
|
||||
processor.logprobs_chunk_size = rows // 2
|
||||
lm_head = Mock(spec=["set_lm_head_pass", "reset_lm_head_pass"])
|
||||
_, _, _, _, chunk_rows = run_with_pool_run_of(96 << 20, False, lm_head)
|
||||
self.assertEqual(chunk_rows, [rows // 2, rows // 2])
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
|
||||
def test_borrow_recovers_from_arena_fragmentation(self):
|
||||
handle = torch.cuda.graph_pool_handle()
|
||||
|
||||
Reference in New Issue
Block a user