Add dummy forward batch preparation hook (#31070)

2026-07-14 14:30:31 -07:00
parent 0d89564d27
commit 08c46e1f1a
2 changed files with 5 additions and 0 deletions
@@ -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.

@@ -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():