[Logprob] Borrow graph-pool memory for input logprob logits construction (#40007)

Co-authored-by: cctry <csycfl@gmail.com>

python/sglang/srt/layers/logits_processor.py

@@ -26,7 +26,7 @@ from sglang.kernels.ops.activation.softcap import (
     softcap_inplace_logits as fused_softcap,
 )
 from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
-from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed import get_attn_tp_group, get_tp_group
 from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
 from sglang.srt.environ import envs
 from sglang.srt.layers import layernorm_sp
@@ -458,7 +458,16 @@ class LogitsProcessor(nn.Module):
             skip_entry_sync=True,
         )
 
-        self.input_logprob_processor = InputLogprobProcessor()
+        chunking_group = None
+        if (
+            self.do_tensor_parallel_all_gather
+            and not self.do_tensor_parallel_all_gather_dp_attn
+        ):
+            group = get_attn_tp_group() if self.use_attn_tp_group else get_tp_group()
+            chunking_group = group.cpu_group
+        self.input_logprob_processor = InputLogprobProcessor(
+            self.vocab_size, chunking_group=chunking_group
+        )
 
     def forward(
         self,

python/sglang/srt/layers/logprob_processor.py

@@ -2,12 +2,20 @@ from __future__ import annotations
 
 import dataclasses
 import logging
+from contextlib import nullcontext
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    is_symmetric_memory_enabled,
+)
 from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_utils.pool import (
+    borrow_graph_pool,
+    graph_pool_borrow_largest_run,
+)
 from sglang.srt.runtime_context import get_exec
 
 if TYPE_CHECKING:
@@ -434,6 +442,10 @@ def _deterministic_inference_enabled() -> bool:
         return False
 
 
+# Scratch allocations and caching-allocator segment rounding in the borrow scope.
+_GRAPH_POOL_BORROW_SLACK_BYTES = 64 << 20
+
+
 class InputLogprobProcessor:
     """Input (prefill) logprob processing: single-pass or chunked.
 
@@ -442,7 +454,13 @@ class InputLogprobProcessor:
     the lm_head / TP-gather machinery in LogitsProcessor.
     """
 
-    def __init__(self):
+    def __init__(
+        self,
+        vocab_size: int,
+        chunking_group: Optional[torch.distributed.ProcessGroup] = None,
+    ):
+        self.vocab_size = vocab_size
+        self.chunking_group = chunking_group
         # enable chunked logprobs processing
         self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
         # chunk size for logprobs processing
@@ -479,6 +497,10 @@ 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)
+
         return self._forward_by_chunk(
             pruned_states,
             sample_indices,
@@ -488,8 +510,37 @@ class InputLogprobProcessor:
             get_logits_fn,
             logits_metadata,
             chunk_size,
+            borrow_logits_memory=borrow_logits_memory,
         )
 
+    def _can_borrow_logits_memory(self, chunk_size: int) -> 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.
+        free_run = (
+            0 if 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:
+            capacity = torch.tensor(fit_rows, dtype=torch.int64, device="cpu")
+            torch.distributed.all_reduce(
+                capacity,
+                op=torch.distributed.ReduceOp.MIN,
+                group=self.chunking_group,
+            )
+            fit_rows = int(capacity.item())
+        # Fall back to the existing reserved workspace if borrowing cannot
+        # hold the whole chunk; do not change LoRA or logprob chunk boundaries.
+        return fit_rows >= chunk_size
+
     def _forward_by_chunk(
         self,
         pruned_states: torch.Tensor,
@@ -500,6 +551,7 @@ class InputLogprobProcessor:
         get_logits_fn: Callable,
         logits_metadata: LogitsMetadata,
         chunk_size: int,
+        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]
@@ -557,14 +609,20 @@ class InputLogprobProcessor:
             # writing through the shared graph logits buffer would alias
             # chunks whose shape happens to match the buffer.
             chunk_states = pruned_states[start_idx:end_idx]
-            chunk_logits = get_logits_fn(
-                chunk_states,
-                lm_head,
-                logits_metadata,
-                use_logits_buffer=num_chunks == 1,
-            )
+            with (
+                borrow_graph_pool(user="input logits")
+                if borrow_logits_memory
+                else nullcontext()
+            ):
+                chunk_logits = get_logits_fn(
+                    chunk_states,
+                    lm_head,
+                    logits_metadata,
+                    use_logits_buffer=num_chunks == 1,
+                )
 
-            # Initialize sampled_logits on first chunk
+            # Sampled outputs must survive graph replay, so they are allocated
+            # outside borrowing. The transient logits are released below.
             if i == 0:
                 sampled_logits = torch.empty(
                     (sample_indices.shape[0], chunk_logits.shape[1]),

python/sglang/srt/managers/scheduler.py

@@ -292,6 +292,7 @@ from sglang.srt.mem_cache.common import (
     retraction_discard,
 )
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.model_executor.runner_utils.pool import prewarm_graph_pool_borrow
 from sglang.srt.model_loader.utils import get_resolved_model_impl
 from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
 from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector
@@ -1134,6 +1135,7 @@ class Scheduler(
             else self.schedule_stream
         )
         with device_module.stream(forward_stream):
+            prewarm_graph_pool_borrow()
             if self.draft_worker is None:
                 model_runner.prewarm_sampling()
             else:

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -174,6 +174,16 @@ def graph_pool_borrow_enabled() -> bool:
     return get_global_graph_memory_pool() is not None
 
 
+def prewarm_graph_pool_borrow() -> None:
+    """Initialize cuBLAS's persistent workspace outside borrowed storage.
+
+    Run on the forward stream before final KV sizing. Otherwise the first
+    borrowed GEMM leaves a cached workspace alive across graph replay.
+    """
+    if graph_pool_borrow_enabled():
+        torch.cuda.current_blas_handle()
+
+
 @contextmanager
 def graph_pool_user_scope(user: str) -> Iterator[None]:
     state = _get_graph_pool_borrow_state()