diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 17bfceb8b..d03288e07 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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, ) diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index 43ca986e8..29805fb3f 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/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()