[Fix] Preserve model runner contracts in prefill CUDA graphs (#35452)

Co-authored-by: Oasis-Git <ayw.sirius19@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -1527,6 +1527,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     column_starts=0,
                     req_lens=shape_inputs["extend_seq_lens"],
                 )
+            forward_batch = self.model_runner.prepare_dummy_forward_batch(forward_batch)
             self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
         return forward_batch, self.model_runner.attn_backend
 
@@ -1860,6 +1861,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # The n-gram hasher runs outside the graph and reads this at replay.
         static_forward_batch.ngram_embedding_info = forward_batch.ngram_embedding_info
         static_forward_batch.engram_history = forward_batch.engram_history
+        static_forward_batch = self.model_runner.prepare_dummy_forward_batch(
+            static_forward_batch
+        )
         if self._is_full_backend:
             forward_batch.next_token_logits_buffer = (
                 static_forward_batch.next_token_logits_buffer
@@ -2070,6 +2074,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
             input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
             input_logprobs_copy_done=output.input_logprobs_copy_done,
+            customized_info=output.customized_info,
             mm_input_embeds=mm_input_embeds,
         )
 

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -8,6 +8,7 @@ import torch
 
 import sglang.srt.model_executor.model_runner_components.cuda_graph_setup as graph_setup
 import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
@@ -140,6 +141,79 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
                 self.assertIs(runner.prepare_dummy_forward_batch(batch), batch)
                 self.assertEqual(batch.attn_tp_sequence_sharded, expected)
 
+    def test_capture_prepare_applies_model_runner_batch_hook(self):
+        class Slot:
+            def __init__(self, buffer):
+                self.buffer = buffer
+
+            def slice_for(self, _bs, num_tokens):
+                return self.buffer[:num_tokens]
+
+        slots = {
+            "input_ids": Slot(torch.zeros(4, dtype=torch.int64)),
+            "out_cache_loc": Slot(torch.zeros(4, dtype=torch.int64)),
+            "positions": Slot(torch.zeros(4, dtype=torch.int64)),
+        }
+        registry = SimpleNamespace(
+            has_slot=lambda name: name in slots,
+            get_slot=lambda name: slots[name],
+        )
+        prepared = []
+        captured = []
+        attention_backend = object()
+        model_runner = SimpleNamespace(
+            model_config=SimpleNamespace(context_len=8),
+            pp_group=SimpleNamespace(is_last_rank=False),
+            ngram_embedding_manager=SimpleNamespace(enabled=False),
+            prepare_dummy_forward_batch=lambda batch: prepared.append(batch) or batch,
+            attn_tp_sequence_sharded=lambda _: False,
+            attn_backend=attention_backend,
+        )
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.model_runner = model_runner
+        runner.device = torch.device("cpu")
+        runner.prefill_backend_name = Backend.BREAKABLE
+        runner.max_context_size = None
+        runner._capture_req_slots = 1
+        runner.max_bs = 4
+        runner._prefill_static_buffers = None
+        runner.buffer_registry = registry
+        runner.require_mlp_tp_gather = False
+        runner.require_attn_tp_gather = False
+        runner._capture_lora = False
+        runner.capture_hidden_mode = CaptureHiddenMode.NULL
+        runner.static_draft_hidden_states = None
+        runner.capture_return_pooled_hidden_states = False
+        runner._next_token_logits_buffer = lambda _rows: None
+        runner._build_capture_spec_info = lambda _num_tokens: None
+        runner._capture_num_token_non_padded = lambda _num_tokens: None
+        runner.tbo_plugin = SimpleNamespace(
+            capture_one_batch_size=lambda batch, **_: captured.append(batch)
+        )
+
+        batch, backend = runner.capture_prepare(4)
+
+        self.assertEqual(prepared, [batch])
+        self.assertEqual(captured, [batch])
+        self.assertIs(backend, attention_backend)
+
+    def test_trim_logits_output_preserves_customized_info(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.model_runner = SimpleNamespace(
+            spec_algorithm=SimpleNamespace(is_speculative=lambda: False),
+        )
+        runner.raw_bs = 1
+        runner._is_full_backend = True
+        customized_info = {"per_request": [object()]}
+        output = runner._trim_logits_output(
+            LogitsProcessorOutput(
+                next_token_logits=torch.zeros((4, 8)),
+                customized_info=customized_info,
+            )
+        )
+
+        self.assertIs(output.customized_info, customized_info)
+
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
         prefill_runner = object()
@@ -247,7 +321,10 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         runner.max_context_size = None
         runner._capture_chunked_prefix = False
         runner.buffer_registry = _FakeBatchRegistry()
-        runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
+        runner.model_runner = SimpleNamespace(
+            attn_tp_sequence_sharded=lambda _: False,
+            prepare_dummy_forward_batch=lambda batch: batch,
+        )
         runner.enable_cp_bcg_capture = False
         runner._is_full_backend = False
         runner.backend = SimpleNamespace()