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

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

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

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.

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",

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