diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 9ebd61ba9..02cedbe4f 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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, diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index 032da6ace..f5144ee37 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/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]), diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 07a1c97cd..8aced8a94 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 5acce066f..fb259065b 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/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() diff --git a/test/registered/unit/layers/test_logprob_chunk_stitching.py b/test/registered/unit/layers/test_logprob_chunk_stitching.py index f84507816..edbfcdfac 100644 --- a/test/registered/unit/layers/test_logprob_chunk_stitching.py +++ b/test/registered/unit/layers/test_logprob_chunk_stitching.py @@ -96,7 +96,7 @@ def _run(proc, batch, chunked, chunk_size): class TestLogprobChunkStitching(CustomTestCase): def _sweep(self, with_token_ids): torch.manual_seed(0) - proc = InputLogprobProcessor() + proc = InputLogprobProcessor(vocab_size=VOCAB) combos = list(coverage_cases(SEQ_SPEC_MENU, max_seqs=4)) self.assertEqual(len(combos), EXPECTED_CASES) tried = 0 diff --git a/test/registered/unit/layers/test_logprob_fast_input.py b/test/registered/unit/layers/test_logprob_fast_input.py index 99e685f4f..3cf285e8f 100644 --- a/test/registered/unit/layers/test_logprob_fast_input.py +++ b/test/registered/unit/layers/test_logprob_fast_input.py @@ -125,7 +125,7 @@ def _shape_of(nested): class TestFastInputLogprobs(CustomTestCase): def _sweep(self, dtype, rtol, atol): torch.manual_seed(0) - proc = InputLogprobProcessor() + proc = InputLogprobProcessor(vocab_size=VOCAB) combos = list(coverage_cases(SEQ_SPEC_MENU, max_seqs=3)) self.assertEqual(len(combos), EXPECTED_CASES) tried = 0 @@ -178,7 +178,7 @@ class TestFastInputLogprobs(CustomTestCase): # rounds at the bf16 logits themselves (normalizer is fp32), so it # sits much closer to the truth than bf16 resolution. torch.manual_seed(0) - proc = InputLogprobProcessor() + proc = InputLogprobProcessor(vocab_size=VOCAB) for combo in coverage_cases(SEQ_SPEC_MENU, max_seqs=3): batch = _build_batch(list(combo), torch.bfloat16) pruned_states, _, input_logprob_indices, _, metadata = batch @@ -233,7 +233,7 @@ class TestFastInputLogprobs(CustomTestCase): # true precision of the result), while the log_softmax path keeps # the logits dtype. Runs on CPU CI so the policy is pinned even # where the CUDA kernels never execute. - proc = InputLogprobProcessor() + proc = InputLogprobProcessor(vocab_size=VOCAB) batch = _build_batch([(4, 1), (3, 0)], torch.bfloat16) got, _ = _run(proc, batch, True, None) self.assertEqual(got.token_logprobs.dtype, torch.float32) @@ -323,7 +323,7 @@ class TestFastInputLogprobs(CustomTestCase): # (the CPU sweeps only cover the torch fallbacks), including the # k > FUSED_TOPK_MAX_K fallback. torch.manual_seed(0) - proc = InputLogprobProcessor() + proc = InputLogprobProcessor(vocab_size=64) for k_override in (None, 20): # k=20 exceeds FUSED_TOPK_MAX_K, exercising the torch fallback; # it needs a vocab that can supply 20 entries. diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index 9a59499b8..06a017cca 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -812,6 +812,10 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): stream=stream_context, Stream=lambda priority: schedule_stream ), ), + patch( + "sglang.srt.managers.scheduler.prewarm_graph_pool_borrow", + side_effect=lambda: trace.append("borrow_prewarm"), + ), self.assertRaisesRegex(RuntimeError, "stop after startup"), ): scheduler.init_model_worker() @@ -830,6 +834,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): "attention", "capture", "stream_enter", + "borrow_prewarm", "prewarm", "stream_exit", "resize", @@ -857,6 +862,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): "attention", "capture", "stream_enter", + "borrow_prewarm", "prewarm", "stream_exit", "resize", @@ -872,6 +878,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase): "attention", "capture", "stream_enter", + "borrow_prewarm", "draft_prewarm", "stream_exit", "resize", 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 1a6d24b1c..59a49400f 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 @@ -251,11 +251,16 @@ class TestGraphPoolBorrow(CustomTestCase): torch.cuda.synchronize() device_id = torch.cuda.current_device() - reserved_before = torch.cuda.memory_reserved(device_id) + borrow_stream = torch.cuda.Stream() with ( envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(True), patch.object(pool, "get_global_graph_memory_pool", return_value=handle), + torch.cuda.stream(borrow_stream), ): + lhs = torch.ones((32, 16), dtype=torch.bfloat16, device="cuda") + copied_product = torch.empty((32, 32), dtype=torch.float32, device="cuda") + pool.prewarm_graph_pool_borrow() + reserved_before = torch.cuda.memory_reserved(device_id) runs = pool.find_free_graph_pool_runs(handle) self.assertGreaterEqual(len(runs), 2) largest_run_bytes = runs[0][1] @@ -301,6 +306,13 @@ class TestGraphPoolBorrow(CustomTestCase): self.assertEqual(reused.data_ptr(), recycled_address) del reused + # The first GEMM on this stream must not cache its workspace + # in borrowed storage, which replay would overwrite. + product = torch.mm(lhs, lhs.T, out_dtype=torch.float32) + self.assertTrue(on_a_run(product)) + copied_product.copy_(product) + del product + # Captures retire the persistent borrow pool. Its storage aliases # existing graph-pool runs, so the reserved footprint is unchanged. pool._teardown_borrow_pool() @@ -315,6 +327,9 @@ class TestGraphPoolBorrow(CustomTestCase): graph.replay() torch.cuda.synchronize() self.assertTrue(torch.equal(y, torch.ones_like(y))) + self.assertTrue( + torch.equal(copied_product, torch.full_like(copied_product, 16)) + ) self.assertEqual(torch.cuda.memory_reserved(device_id), reserved_before) del graph, y