From 08c46e1f1af0f389b32bc4cd58da9e528889de9d Mon Sep 17 00:00:00 2001 From: paulzhang-tm Date: Tue, 14 Jul 2026 14:30:31 -0700 Subject: [PATCH] Add dummy forward batch preparation hook (#31070) --- python/sglang/srt/model_executor/model_runner.py | 4 ++++ python/sglang/srt/model_executor/runner/base_runner.py | 1 + 2 files changed, 5 insertions(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 461705502..cb81353c2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1031,6 +1031,10 @@ class ModelRunner: def update_decode_attn_backend(self, stream_idx: int): self.decode_attn_backend = self.decode_attn_backend_group[stream_idx] + def prepare_dummy_forward_batch(self, forward_batch: ForwardBatch) -> ForwardBatch: + """Customize a runner-created dummy batch before attention metadata initialization.""" + return forward_batch + def _prepare_eager_forward_batch(self, forward_batch: ForwardBatch) -> None: """Pad / normalize a batch for the eager (non-cuda-graph) forward. diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 5b8727c69..e6e8b8a4a 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -552,6 +552,7 @@ class BaseRunner(ABC): if lora_ids is not None: mr.lora_manager.prepare_lora_batch(forward_batch) + forward_batch = mr.prepare_dummy_forward_batch(forward_batch) mr.attn_backend.init_forward_metadata(forward_batch) def run_once():