[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.