From 9e5a62a767433b66a0e608276ca520d4e53e765b Mon Sep 17 00:00:00 2001 From: metamergebot Date: Sat, 19 Sep 2026 12:03:52 -0700 Subject: [PATCH] [Logprob] Serve input-logprob temporaries from CUDA-graph-pool dead space (#40038) Co-authored-by: cctry --- python/sglang/srt/layers/logits_processor.py | 23 ++ python/sglang/srt/layers/logprob_processor.py | 344 +++++++++++------- .../batch_result_processor.py | 1 + python/sglang/srt/managers/utils.py | 17 +- .../runner/prefill_cuda_graph_runner.py | 1 + python/sglang/srt/utils/common.py | 10 + .../layers/test_logprob_memory_distributed.py | 163 +++++++++ .../layers/test_logprob_chunk_stitching.py | 22 ++ .../unit/layers/test_logprob_fast_input.py | 2 + .../runner_utils/test_graph_pool_borrow.py | 241 ++++++++++++ 10 files changed, 674 insertions(+), 150 deletions(-) create mode 100644 test/registered/e2e/layers/test_logprob_memory_distributed.py diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 3c50a8475..f3c97fc2b 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -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( diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index f5144ee37..f2dcf7d7a 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/python/sglang/srt/layers/logprob_processor.py @@ -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, ) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 386e8d10f..7ce4572f5 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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() diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index bfcb97124..322dba9ef 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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 { diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 5ffbf697d..bcee09dbe 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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, ) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 5802c11b1..22ed8e168 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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" diff --git a/test/registered/e2e/layers/test_logprob_memory_distributed.py b/test/registered/e2e/layers/test_logprob_memory_distributed.py new file mode 100644 index 000000000..cbeacc64b --- /dev/null +++ b/test/registered/e2e/layers/test_logprob_memory_distributed.py @@ -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"])) diff --git a/test/registered/unit/layers/test_logprob_chunk_stitching.py b/test/registered/unit/layers/test_logprob_chunk_stitching.py index edbfcdfac..51e5cafd9 100644 --- a/test/registered/unit/layers/test_logprob_chunk_stitching.py +++ b/test/registered/unit/layers/test_logprob_chunk_stitching.py @@ -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() diff --git a/test/registered/unit/layers/test_logprob_fast_input.py b/test/registered/unit/layers/test_logprob_fast_input.py index 3cf285e8f..ce72b759b 100644 --- a/test/registered/unit/layers/test_logprob_fast_input.py +++ b/test/registered/unit/layers/test_logprob_fast_input.py @@ -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))], diff --git a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py index 831bf434a..2d49cb46f 100644 --- a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -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()