[Memory] Size the CUDA graph pool from warmup measurements and fix graph-pool borrowing (#36911)

Co-authored-by: cctry <cctry@fb.com>

python/sglang/srt/environ.py

@@ -1293,6 +1293,9 @@ class Envs:
     SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False)
     # Reuse wholly-free graph-pool segments for step-local eager allocations.
     SGLANG_ENABLE_GRAPH_POOL_BORROW = EnvBool(False)
+    # Mint capture's measured footprint as one span so the graph pool is carved
+    # out of a single contiguous region instead of grown segment by segment.
+    SGLANG_ENABLE_GRAPH_POOL_PRECARVE = EnvBool(False)
     # Eager forward wraps the ForwardBatch's own tensors instead of copying them
     # into the CUDA graph buffer registry (no per-iter device-to-device copy).
     SGLANG_EAGER_INPUT_NO_COPY = EnvBool(False)

python/sglang/srt/managers/scheduler.py

@@ -1060,6 +1060,13 @@ class Scheduler(
         self.init_all_cuda_graphs()
 
         model_runner = self.tp_worker.model_runner
+        with torch.get_device_module(model_runner.device).stream(
+            model_runner.forward_stream
+        ):
+            if self.draft_worker is None:
+                model_runner.prewarm_sampling()
+            else:
+                self.draft_worker.prewarm_sampling()
         if model_runner.token_to_kv_pool.post_capture_active:
             tic = time.perf_counter()
             model_runner.post_capture_resize_kv_pool()

python/sglang/srt/model_executor/model_runner.py

@@ -253,6 +253,14 @@ elif current_platform.is_out_of_tree():
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class SamplingPrewarmResult:
+    """Memory requirements observed while pre-warming a sampling path."""
+
+    sampling_input_bytes: int = 0
+    sampling_headroom_bytes: int = 0
+
+
 def _prefill_cuda_graph_allows_context_parallel(
     prefill_runner, forward_batch: ForwardBatch
 ) -> bool:
@@ -377,6 +385,7 @@ class ModelRunner:
         self.draft_model_idx = draft_model_idx
         self.enable_hisparse = get_memory().enable_hisparse
         self._sampling_observer: Optional[SamplingObserver] = None
+        self.sampling_prewarm_result = SamplingPrewarmResult()
 
         self.init_startup_observability()
 
@@ -1058,6 +1067,11 @@ class ModelRunner:
             n_prepared,
         )
 
+    def prewarm_sampling(self) -> SamplingPrewarmResult:
+        """Warm the sampling path after graph initialization."""
+        self.sampling_prewarm_result = SamplingPrewarmResult()
+        return self.sampling_prewarm_result
+
     def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True):
         capture = capture_cuda_graphs(
             model_runner=self, capture_decode_cuda_graph=capture_decode_cuda_graph

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -11,6 +11,7 @@ from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.distributed import get_world_group
 from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb
 from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_utils.pool import graph_pool_borrow_enabled
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_disagg,
@@ -85,6 +86,13 @@ def compute_post_capture_kv_resize(
             )
             / 1024,
         )
+    if not graph_pool_borrow_enabled():
+        # Borrowing serves the sampling temporaries out of idle graph storage;
+        # without it they need real headroom the KV pool must not claim.
+        headroom_gb = max(
+            headroom_gb,
+            model_runner.sampling_prewarm_result.sampling_headroom_bytes / (1 << 30),
+        )
     mm_reservation_gb = mm_runtime_reservation_gb(
         is_multimodal=model_runner.model_config.is_multimodal,
         mm_feature_transport=get_mm().mm_feature_transport,

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -40,6 +40,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
     enable_breakable_cuda_graph,
 )
 from sglang.srt.model_executor.runner_utils.pool import (
+    GraphPoolPrecarve,
     get_or_create_global_graph_memory_pool,
     graph_pool_capture_scope,
     graph_pool_replay_scope,
@@ -77,6 +78,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         self._capture_stream: Optional[torch.cuda.Stream] = None
         self._debug_eager = debug_eager
         self._shared_output_buffer: Optional[Any] = None
+        self._precarve = GraphPoolPrecarve()
         self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
             and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
@@ -117,7 +119,8 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         for _ in range(2):
             self._device_module.synchronize()
             self._tp_group.barrier()
-            warmup_out = forward_fn()
+            with self._precarve.measure():
+                warmup_out = forward_fn()
             if post_warmup_hook is not None:
                 post_warmup_hook()
 
@@ -134,6 +137,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
             stream=self._capture_stream,
             barrier_fn=self._tp_group.barrier,
         ):
+            self._precarve.mint()
             out = captured_fn()
             out_rows = self._output_rows(out, size)
             self._copy_output_to_buffer(out, self._shared_output_buffer, out_rows)

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -31,6 +31,7 @@ from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
 from sglang.srt.model_executor.runner_utils.pool import (
+    GraphPoolPrecarve,
     get_or_create_global_graph_memory_pool,
     graph_pool_capture_scope,
     graph_pool_replay_scope,
@@ -64,6 +65,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         self._device_module = cuda_graph_runner.device_module
         self._tp_group = cuda_graph_runner.model_runner.tp_group
         self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._precarve = GraphPoolPrecarve()
         self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
             and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
@@ -107,7 +109,8 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         for _ in range(2):
             self._device_module.synchronize()
             self._tp_group.barrier()
-            forward_fn()
+            with self._precarve.measure():
+                forward_fn()
             if profiler is not None:
                 profiler.step()
             if post_warmup_hook is not None:
@@ -131,6 +134,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
             graph_pool_capture_scope(),
             graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream),
         ):
+            self._precarve.mint()
             out = forward_fn()
 
         if profiler is not None:

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

@@ -46,6 +46,7 @@ def disable_graph_pool_borrow(reason: str) -> None:
     """Disable borrowing when graph storage is managed outside the shared pool."""
     global _borrow_disabled_reason
     _borrow_disabled_reason = reason
+    _teardown_borrow_pool()
     logger.info("Graph pool borrow disabled: %s", reason)
 
 
@@ -93,6 +94,43 @@ def get_or_create_global_graph_capture_stream() -> Any:
     return get_stream(_CAPTURE_STREAM_NAME)
 
 
+class GraphPoolPrecarve:
+    """Pre-carve the memory pool to reduce fragmentation."""
+
+    def __init__(self) -> None:
+        self.nbytes = 0
+        self.minted = False
+
+    @contextmanager
+    def measure(self) -> Iterator[None]:
+        """Wrap one eager warmup. the last one before ``mint`` sets the size."""
+        if self.minted or not envs.SGLANG_ENABLE_GRAPH_POOL_PRECARVE.get():
+            yield
+            return
+        torch.cuda.synchronize()
+        # Shrink the cache first so the warmup's reserved growth is its own
+        # footprint. Reserved (not allocated) is the stat to use: allocated
+        # peak is the live-byte sum and undershoots by exactly the packing
+        # holes the carved span has to absorb.
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        base = torch.cuda.memory_stats()["reserved_bytes.all.current"]
+        yield
+        torch.cuda.synchronize()
+        self.nbytes = torch.cuda.memory_stats()["reserved_bytes.all.peak"] - base
+
+    def mint(self) -> None:
+        """Pre-allocate the space"""
+        if self.minted:
+            return
+        self.minted = True
+        if self.nbytes <= 0:
+            return
+        span = torch.empty(self.nbytes, dtype=torch.uint8, device="cuda")
+        del span
+        logger.info("Graph pool pre-carved: %.2f GB", self.nbytes / 2**30)
+
+
 def graph_pool_borrow_enabled() -> bool:
     if (
         _borrow_disabled_reason is not None
@@ -143,11 +181,11 @@ def graph_pool_capture_scope() -> Iterator[None]:
 
 
 def find_free_graph_pool_runs(pool_id: Any) -> list[tuple[int, int]]:
-    """Return contiguous inactive runs inside wholly-free segments, largest first."""
+    """Return contiguous inactive runs across all pool segments, largest first;
+    live and pending-free blocks break runs, so a run never overlaps live data.
+    """
     runs: list[tuple[int, int]] = []
     for segment in torch.cuda.memory_snapshot(pool_id, include_traces=False):
-        if segment["allocated_size"] != 0:
-            continue
         run_address = 0
         run_bytes = 0
         for block in segment["blocks"]:
@@ -195,16 +233,18 @@ def borrow_graph_pool(user: str) -> Iterator[None]:
             # Return completed cross-stream frees to the cache. The allocator
             # processes their events on a later allocation.
             torch.empty(1, device="cuda")
-            if (
-                _borrow_stub.freed_bytes
-                or _borrow_stub.cursor_bytes > _borrow_extents_total // 2
-            ):
-                # Rebuild if empty_cache released the arena's segments or if
-                # unresolved deferred frees consumed half of the extents.
+            if _borrow_stub.freed_bytes:
+                # Rebuild if the caching allocator released an arena segment.
+                # A high cursor alone means the cache owns reusable segments;
+                # rebuilding discards them and can turn the next large borrow
+                # into a fragmented cold-allocation OOM.
                 _teardown_borrow_pool()
         if _borrow_mem_pool is None:
             if _borrow_stub is None:
                 _borrow_stub = BumpArenaStub()
+                # Runs are sorted largest first, so first fit would carve every
+                # small allocation out of the run a probability matrix needs.
+                _borrow_stub.set_best_fit(True)
             if _borrow_static_runs is not None:
                 runs = _borrow_static_runs
             else:

python/sglang/srt/speculative/base_spec_worker.py

@@ -19,7 +19,10 @@ if TYPE_CHECKING:
         UpdateWeightsFromIPCReqInput,
     )
     from sglang.srt.managers.tp_worker import TpModelWorker
-    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.model_runner import (
+        ModelRunner,
+        SamplingPrewarmResult,
+    )
     from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
 
@@ -215,6 +218,9 @@ class BaseSpecWorker(ABC):
             return 0.0
         return self.draft_worker.weight_load_time
 
+    def prewarm_sampling(self) -> SamplingPrewarmResult:
+        return self.target_worker.model_runner.prewarm_sampling()
+
     @property
     def preloaded_weights_bytes(self) -> int:
         if self.draft_worker is None:

python/sglang/srt/speculative/dflash_utils.py

@@ -17,6 +17,7 @@ from sglang.srt.layers.sampler import (
     top_p_normalize_probs_torch,
 )
 from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.model_executor.runner_utils.pool import borrow_graph_pool
 from sglang.srt.runtime_context import get_spec
 from sglang.srt.speculative.spec_utils import sample_simulated_acc_len
 from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
@@ -924,44 +925,22 @@ def compute_dflash_sampling_correct_drafts_and_bonus(
 
     device = next_token_logits.device
 
-    if uniform_samples is None:
-        uniform_samples = torch.rand(
-            (bs, draft_token_num), dtype=torch.float32, device=device
+    if uniform_samples is not None and uniform_samples.shape != (bs, draft_token_num):
+        raise ValueError(
+            "uniform_samples shape mismatch. "
+            f"Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}."
         )
-    else:
-        if uniform_samples.shape != (bs, draft_token_num):
-            raise ValueError(
-                "uniform_samples shape mismatch. "
-                f"Expected {(bs, draft_token_num)}, got {tuple(uniform_samples.shape)}."
-            )
-        uniform_samples = uniform_samples.to(device=device, dtype=torch.float32)
-
-    if uniform_samples_for_final_sampling is None:
-        uniform_samples_for_final_sampling = torch.rand(
-            (bs,), dtype=torch.float32, device=device
-        )
-    else:
-        if uniform_samples_for_final_sampling.shape != (bs,):
-            raise ValueError(
-                "uniform_samples_for_final_sampling shape mismatch. "
-                f"Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}."
-            )
-        uniform_samples_for_final_sampling = uniform_samples_for_final_sampling.to(
-            device=device,
-            dtype=torch.float32,
+    if (
+        uniform_samples_for_final_sampling is not None
+        and uniform_samples_for_final_sampling.shape != (bs,)
+    ):
+        raise ValueError(
+            "uniform_samples_for_final_sampling shape mismatch. "
+            f"Expected {(bs,)}, got {tuple(uniform_samples_for_final_sampling.shape)}."
         )
 
-    target_probs = build_dflash_verify_target_probs(
-        next_token_logits=next_token_logits,
-        sampling_info=sampling_info,
-        draft_token_num=draft_token_num,
-        bs=bs,
-        max_top_k=max_top_k,
-        uniform_top_k_value=uniform_top_k_value,
-        use_sparse_topk=use_sparse_topk,
-    )
-    draft_probs = torch.zeros_like(target_probs)
-
+    # Cached across steps, and `correct_len` below aliases `accept_token_num`,
+    # so these must predate the borrow scope the next replay reclaims.
     (
         retrieve_index,
         retrieve_next_token,
@@ -974,25 +953,59 @@ def compute_dflash_sampling_correct_drafts_and_bonus(
         draft_token_num=draft_token_num,
         device=device,
     )
-    candidates_i64 = (
-        candidates if candidates.dtype == torch.int64 else candidates.to(torch.int64)
-    )
-    tree_speculative_sampling_target_only(
-        predicts=predicts,
-        accept_index=accept_index,
-        accept_token_num=accept_token_num,
-        candidates=candidates_i64,
-        retrive_index=retrieve_index,
-        retrive_next_token=retrieve_next_token,
-        retrive_next_sibling=retrieve_next_sibling,
-        uniform_samples=uniform_samples,
-        uniform_samples_for_final_sampling=uniform_samples_for_final_sampling,
-        target_probs=target_probs,
-        draft_probs=draft_probs,
-        threshold_single=threshold_single,
-        threshold_acc=threshold_acc,
-        deterministic=True,
-    )
+
+    # The full-vocabulary matrices die with this step, so their bytes may come
+    # from the graph pool's idle storage. Anything outliving the scope is not.
+    with borrow_graph_pool(user="DFLASH verify probabilities"):
+        if uniform_samples is None:
+            coins = torch.rand(
+                (bs, draft_token_num), dtype=torch.float32, device=device
+            )
+        else:
+            coins = uniform_samples.to(device=device, dtype=torch.float32)
+        if uniform_samples_for_final_sampling is None:
+            coins_for_final_sampling = torch.rand(
+                (bs,), dtype=torch.float32, device=device
+            )
+        else:
+            coins_for_final_sampling = uniform_samples_for_final_sampling.to(
+                device=device,
+                dtype=torch.float32,
+            )
+
+        target_probs = build_dflash_verify_target_probs(
+            next_token_logits=next_token_logits,
+            sampling_info=sampling_info,
+            draft_token_num=draft_token_num,
+            bs=bs,
+            max_top_k=max_top_k,
+            uniform_top_k_value=uniform_top_k_value,
+            use_sparse_topk=use_sparse_topk,
+        )
+        draft_probs = torch.zeros_like(target_probs)
+        candidates_i64 = (
+            candidates
+            if candidates.dtype == torch.int64
+            else candidates.to(torch.int64)
+        )
+        tree_speculative_sampling_target_only(
+            predicts=predicts,
+            accept_index=accept_index,
+            accept_token_num=accept_token_num,
+            candidates=candidates_i64,
+            retrive_index=retrieve_index,
+            retrive_next_token=retrieve_next_token,
+            retrive_next_sibling=retrieve_next_sibling,
+            uniform_samples=coins,
+            uniform_samples_for_final_sampling=coins_for_final_sampling,
+            target_probs=target_probs,
+            draft_probs=draft_probs,
+            threshold_single=threshold_single,
+            threshold_acc=threshold_acc,
+            deterministic=True,
+        )
+        del target_probs, draft_probs, candidates_i64
+        del coins, coins_for_final_sampling
 
     correct_len = accept_token_num
     row_ids = torch.arange(bs, dtype=torch.long, device=device)

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -33,17 +33,24 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardMode,
     compute_position,
 )
+from sglang.srt.model_executor.model_runner import SamplingPrewarmResult
+from sglang.srt.model_executor.runner_utils.pool import (
+    disable_graph_pool_borrow,
+    graph_pool_borrow_enabled,
+)
 from sglang.srt.runtime_context import (
     get_exec,
     get_schedule,
     get_spec,
     mamba_track_grid,
 )
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
 from sglang.srt.speculative.dflash_info import DFlashVerifyInput
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
 from sglang.srt.speculative.dflash_utils import (
+    _get_or_create_chain_verify_buffers,
     apply_dflash_simulated_acceptance,
     apply_dflash_verify_logits_adjustments,
     can_dflash_use_fused_qkv_proj,
@@ -491,6 +498,119 @@ class DFlashWorkerV2(BaseSpecWorker):
             capture_decode_cuda_graph=capture_decode_cuda_graph
         )
 
+    def _prewarm_batch_size(self, block_size: int) -> int:
+        """Largest batch the non-greedy verify path can see in one step."""
+        bs = self.model_runner.max_decode_logits_rows() // block_size
+        max_running_requests = self.model_runner.max_running_requests
+        if max_running_requests is not None:
+            bs = min(bs, int(max_running_requests))
+        return max(1, bs)
+
+    def _measure_sampling_peak(
+        self,
+        *,
+        candidates: torch.Tensor,
+        target_logits: torch.Tensor,
+        sampling_info: SamplingBatchInfo,
+    ) -> int:
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        base_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]
+        compute_dflash_sampling_correct_drafts_and_bonus(
+            candidates=candidates,
+            next_token_logits=target_logits,
+            sampling_info=sampling_info,
+            max_top_k=None,
+            uniform_top_k_value=None,
+        )
+        torch.cuda.synchronize()
+        return torch.cuda.memory_stats()["allocated_bytes.all.peak"] - base_bytes
+
+    def prewarm_sampling(self) -> SamplingPrewarmResult:
+        """Rehearse the non-greedy verify path so an undersized graph-pool
+        borrow surfaces at startup rather than under traffic."""
+        if not is_cuda() or not is_dflash_sampling_verify_available():
+            return self.model_runner.prewarm_sampling()
+
+        block_size = int(self.block_size)
+        vocab_size = self.model_config.vocab_size
+        bs = self._prewarm_batch_size(block_size)
+        device = self.device
+
+        sampling_info = SamplingBatchInfo(
+            temperatures=torch.ones((bs, 1), dtype=torch.float32, device=device),
+            top_ps=torch.full((bs,), 0.95, dtype=torch.float32, device=device),
+            top_ks=torch.zeros(bs, dtype=torch.int32, device=device),
+            min_ps=torch.zeros(bs, dtype=torch.float32, device=device),
+            is_all_greedy=False,
+            is_any_greedy=False,
+            need_top_p_sampling=True,
+            # Top-k off: the sparse path is cheaper, so the dense full-vocab
+            # renormalization is the peak this rehearsal has to reproduce.
+            need_top_k_sampling=False,
+            need_min_p_sampling=False,
+            vocab_size=vocab_size,
+        )
+        candidates = torch.zeros((bs, block_size), dtype=torch.int64, device=device)
+        target_logits = torch.zeros(
+            (bs * block_size, vocab_size), dtype=torch.float32, device=device
+        )
+
+        # Size the process-lifetime chain buffers before anything borrows. The
+        # cache is keyed on device index, so pass the resolved device, not "cuda".
+        _get_or_create_chain_verify_buffers(
+            bs=bs, draft_token_num=block_size, device=target_logits.device
+        )
+        # One unborrowed pass, so any workspace a kernel caches for the process
+        # lifetime lands outside graph storage and off the measured peak.
+        with envs.SGLANG_ENABLE_GRAPH_POOL_BORROW.override(False):
+            self._measure_sampling_peak(
+                candidates=candidates,
+                target_logits=target_logits,
+                sampling_info=sampling_info,
+            )
+
+        borrowing = graph_pool_borrow_enabled()
+        try:
+            peak_bytes = self._measure_sampling_peak(
+                candidates=candidates,
+                target_logits=target_logits,
+                sampling_info=sampling_info,
+            )
+        except torch.OutOfMemoryError:
+            if not borrowing:
+                raise
+            reason = (
+                "DFLASH sampling rehearsal exhausted graph-pool memory for the "
+                f"{bs}x{block_size}x{vocab_size} verify probability matrices"
+            )
+            logger.warning(
+                "Graph pool %s; disabling borrowing and reserving the measured "
+                "headroom in the post-capture KV sizing instead",
+                reason,
+            )
+            disable_graph_pool_borrow(reason)
+            peak_bytes = self._measure_sampling_peak(
+                candidates=candidates,
+                target_logits=target_logits,
+                sampling_info=sampling_info,
+            )
+
+        sampling_input_bytes = target_logits.numel() * target_logits.element_size()
+        result = SamplingPrewarmResult(
+            sampling_input_bytes=sampling_input_bytes,
+            sampling_headroom_bytes=peak_bytes + sampling_input_bytes,
+        )
+        self.model_runner.sampling_prewarm_result = result
+        logger.info(
+            "DFLASH sampling pre-warm: borrowing=%s bs=%d peak=%d B headroom=%d B",
+            graph_pool_borrow_enabled(),
+            bs,
+            peak_bytes,
+            result.sampling_headroom_bytes,
+        )
+        return result
+
     def _maybe_build_draft_sampler(self):
         def _eager(reason):
             if self.ps.tp_rank == 0:

python/sglang/srt/speculative/eagle_utils.py

@@ -19,7 +19,6 @@ from sglang.srt.mem_cache.allocation_sizing import (
     get_alloc_reserve_per_decode,
     page_aligned_decode_alloc_lens,
 )
-from sglang.srt.model_executor.runner_utils.pool import borrow_graph_pool
 from sglang.srt.runtime_context import get_parallel, get_spec
 from sglang.srt.utils import (
     is_cpu,
@@ -789,86 +788,78 @@ def eagle_sample(
             else tree_speculative_sampling_target_only
         )
 
-        # These full-vocabulary matrices are consumed by the sampling kernel
-        # within this step. Returned tensors were allocated before the scope,
-        # so the next CUDA graph replay may safely reclaim these borrowed bytes.
-        with borrow_graph_pool(user="EAGLE probability borrow"):
-            expanded_temperature = torch.repeat_interleave(
-                sampling_info.temperatures, verify_input.draft_token_num, dim=0
-            )  # (bs * num_draft_tokens, 1)
+        expanded_temperature = torch.repeat_interleave(
+            sampling_info.temperatures, verify_input.draft_token_num, dim=0
+        )  # (bs * num_draft_tokens, 1)
 
-            target_probs = F.softmax(
-                next_token_logits / expanded_temperature, dim=-1
-            )  # (bs * num_draft_tokens, vocab_size)
-            maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
-            if sampling_info.need_top_k_sampling:
-                target_probs = top_k_renorm_prob(
-                    target_probs,
-                    torch.repeat_interleave(
-                        sampling_info.top_ks, verify_input.draft_token_num, dim=0
-                    ),
-                )  # (bs * num_draft_tokens, vocab_size)
-                maybe_detect_nan(
-                    target_probs, "v2 verify: target_probs after top_k_renorm"
-                )
-            if sampling_info.need_top_p_sampling:
-                target_probs = top_p_renorm_prob(
-                    target_probs,
-                    torch.repeat_interleave(
-                        sampling_info.top_ps, verify_input.draft_token_num, dim=0
-                    ),
-                )
-                maybe_detect_nan(
-                    target_probs, "v2 verify: target_probs after top_p_renorm"
-                )
-            target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1)
-            draft_probs = (
-                verify_input.draft_probs
-                if use_rejection_sampling
-                else torch.zeros_like(target_probs)
-            )
-            # Defense-in-depth behind the spec_hook startup allowlist: validate
-            # the actual kernel inputs before the Triton kernel.
-            if use_rejection_sampling and (
-                draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1]
-            ):
-                raise ValueError(
-                    "Rejection sampling requires a target-vocab draft proposal "
-                    "distribution; the current speculative algorithm/draft worker "
-                    "does not produce one (draft_probs missing or vocab-mismatched)."
-                )
-
-            coins, coins_for_final_sampling = _verify_coins(
-                sampling_info=sampling_info,
-                seq_lens=batch.seq_lens,
-                draft_token_num=verify_input.draft_token_num,
-                candidates=candidates,
-                device=device,
-            )
-            sampling_fn(
-                predicts=predict,  # mutable
-                accept_index=accept_index,  # mutable
-                accept_token_num=num_correct_drafts,  # mutable
-                candidates=candidates,
-                # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
-                retrive_index=verify_input.retrieve_index,
-                retrive_next_token=verify_input.retrieve_next_token,
-                retrive_next_sibling=verify_input.retrieve_next_sibling,
-                uniform_samples=coins,
-                uniform_samples_for_final_sampling=coins_for_final_sampling,
-                target_probs=target_probs,
-                draft_probs=draft_probs,
-                threshold_single=get_spec().speculative_accept_threshold_single,
-                threshold_acc=get_spec().speculative_accept_threshold_acc,
-                deterministic=True,
-            )
-            del (
-                expanded_temperature,
+        target_probs = F.softmax(
+            next_token_logits / expanded_temperature, dim=-1
+        )  # (bs * num_draft_tokens, vocab_size)
+        maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
+        if sampling_info.need_top_k_sampling:
+            target_probs = top_k_renorm_prob(
                 target_probs,
-                draft_probs,
-                coins,
-                coins_for_final_sampling,
+                torch.repeat_interleave(
+                    sampling_info.top_ks, verify_input.draft_token_num, dim=0
+                ),
+            )  # (bs * num_draft_tokens, vocab_size)
+            maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
+        if sampling_info.need_top_p_sampling:
+            target_probs = top_p_renorm_prob(
+                target_probs,
+                torch.repeat_interleave(
+                    sampling_info.top_ps, verify_input.draft_token_num, dim=0
+                ),
             )
+            maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm")
+        target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1)
+        draft_probs = (
+            verify_input.draft_probs
+            if use_rejection_sampling
+            else torch.zeros_like(target_probs)
+        )
+        # Defense-in-depth behind the spec_hook startup allowlist: validate
+        # the actual kernel inputs before the Triton kernel.
+        if use_rejection_sampling and (
+            draft_probs is None or draft_probs.shape[-1] != target_probs.shape[-1]
+        ):
+            raise ValueError(
+                "Rejection sampling requires a target-vocab draft proposal "
+                "distribution; the current speculative algorithm/draft worker "
+                "does not produce one (draft_probs missing or vocab-mismatched)."
+            )
+
+        coins, coins_for_final_sampling = _verify_coins(
+            sampling_info=sampling_info,
+            seq_lens=batch.seq_lens,
+            draft_token_num=verify_input.draft_token_num,
+            candidates=candidates,
+            device=device,
+        )
+        sampling_fn(
+            predicts=predict,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=num_correct_drafts,  # mutable
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
+            retrive_index=verify_input.retrieve_index,
+            retrive_next_token=verify_input.retrieve_next_token,
+            retrive_next_sibling=verify_input.retrieve_next_sibling,
+            uniform_samples=coins,
+            uniform_samples_for_final_sampling=coins_for_final_sampling,
+            target_probs=target_probs,
+            draft_probs=draft_probs,
+            threshold_single=get_spec().speculative_accept_threshold_single,
+            threshold_acc=get_spec().speculative_accept_threshold_acc,
+            deterministic=True,
+        )
+        del (
+            expanded_temperature,
+            target_probs,
+            draft_probs,
+            coins,
+            coins_for_final_sampling,
+        )
 
         # Sync sampling results across TP ranks: different GPUs may
         # produce slightly different target_probs due to floating-point

python/sglang/srt/utils/cuda_vmm_utils.py

@@ -358,6 +358,7 @@ static size_t g_cursors[BUMPARENA_MAX_EXTENTS];
 static size_t g_num_extents = 0;
 static size_t g_freed_bytes = 0;
 static size_t g_align = 512;
+static int g_best_fit = 0;
 static std::mutex g_mu;
 static size_t align_up(size_t v, size_t a){{ return (v + a - 1) / a * a; }}
 void bumparena_set_extents_{sfx}(const uintptr_t* bases, const size_t* sizes, size_t n){{
@@ -372,6 +373,7 @@ void bumparena_set_extents_{sfx}(const uintptr_t* bases, const size_t* sizes, si
   }}
 }}
 void bumparena_set_align_{sfx}(size_t a){{ std::lock_guard<std::mutex> lk(g_mu); if (a) g_align=a; }}
+void bumparena_set_best_fit_{sfx}(int on){{ std::lock_guard<std::mutex> lk(g_mu); g_best_fit = on; }}
 size_t bumparena_cursor_{sfx}(void){{
   std::lock_guard<std::mutex> lk(g_mu);
   size_t total = 0;
@@ -380,14 +382,21 @@ size_t bumparena_cursor_{sfx}(void){{
 }}
 void* bumparena_malloc_{sfx}(size_t size, int device, void* stream){{
   std::lock_guard<std::mutex> lk(g_mu);
+  size_t need = align_up(size, g_align);
+  size_t pick = BUMPARENA_MAX_EXTENTS;
   for (size_t i = 0; i < g_num_extents; ++i) {{
-    size_t need = g_cursors[i] + align_up(size, g_align);
-    if (need > g_reserved[i]) continue;
-    void* p = reinterpret_cast<void*>(g_bases[i] + g_cursors[i]);
-    g_cursors[i] = need;
-    return p;
+    size_t avail = g_reserved[i] - g_cursors[i];
+    if (avail < need) continue;
+    // First fit is for callers that map physical pages at the offsets they
+    // are handed, so extent order is theirs to choose.
+    if (!g_best_fit) {{ pick = i; break; }}
+    if (pick == BUMPARENA_MAX_EXTENTS ||
+        avail < g_reserved[pick] - g_cursors[pick]) pick = i;
   }}
-  return 0;   // no extent fits -- surfaces as an allocator OOM
+  if (pick == BUMPARENA_MAX_EXTENTS) return 0;   // no extent fits -- surfaces as an allocator OOM
+  void* p = reinterpret_cast<void*>(g_bases[pick] + g_cursors[pick]);
+  g_cursors[pick] += need;
+  return p;
 }}
 size_t bumparena_freed_{sfx}(void){{ std::lock_guard<std::mutex> lk(g_mu); return g_freed_bytes; }}
 void bumparena_free_{sfx}(void* ptr, size_t size, int device, void* stream){{
@@ -401,8 +410,8 @@ void bumparena_free_{sfx}(void* ptr, size_t size, int device, void* stream){{
 class BumpArenaStub:
     """JIT-built pluggable bump allocator over caller-provided device VA extents.
 
-    ``malloc`` first-fits an extent and hands out ``base + cursor``; ``free``
-    is a no-op. Plain ``torch.empty`` can thus be placed on externally managed
+    ``malloc`` picks an extent and hands out ``base + cursor``; ``free`` is a
+    no-op. Plain ``torch.empty`` can thus be placed on externally managed
     storage by wrapping ``allocator`` in a ``torch.cuda.MemPool``.
     ``set_extents`` re-points the arena and resets every cursor, letting one
     stub serve successive region sets.
@@ -462,6 +471,9 @@ class BumpArenaStub:
         self._fn_set_align = lib[f"bumparena_set_align_{self.sfx}"]
         self._fn_set_align.argtypes = [ctypes.c_size_t]
         self._fn_set_align.restype = None
+        self._fn_set_best_fit = lib[f"bumparena_set_best_fit_{self.sfx}"]
+        self._fn_set_best_fit.argtypes = [ctypes.c_int]
+        self._fn_set_best_fit.restype = None
         self._fn_cursor = lib[f"bumparena_cursor_{self.sfx}"]
         self._fn_cursor.argtypes = []
         self._fn_cursor.restype = ctypes.c_size_t
@@ -486,6 +498,9 @@ class BumpArenaStub:
     def set_align(self, nbytes: int) -> None:
         self._fn_set_align(ctypes.c_size_t(nbytes))
 
+    def set_best_fit(self, on: bool) -> None:
+        self._fn_set_best_fit(ctypes.c_int(1 if on else 0))
+
     @property
     def cursor_bytes(self) -> int:
         return int(self._fn_cursor())

python/sglang/test/kits/eval_accuracy_kit.py

@@ -194,6 +194,10 @@ class GSM8KMixin:
     gsm8k_thinking: bool = False  # sgl_eval backend
     gsm8k_max_tokens: Optional[int] = None  # sgl_eval backend
     gsm8k_n_repeats: int = 1  # sgl_eval backend
+    # None keeps run_eval's greedy default; set both to route the run through
+    # the sampling path.
+    gsm8k_temperature: Optional[float] = None
+    gsm8k_top_p: Optional[float] = None
 
     def test_gsm8k(self):
         requests.get(self.base_url + "/flush_cache")
@@ -228,6 +232,8 @@ class GSM8KMixin:
                 api="completion",
                 max_tokens=512,
                 num_shots=self.gsm8k_num_shots,
+                temperature=self.gsm8k_temperature,
+                top_p=self.gsm8k_top_p,
             )