From b3a0185cab0411bdfe26a2ae9f3485499e03d6cd Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Sat, 18 Jul 2026 12:57:54 -0400 Subject: [PATCH] model_runner: extract post-memory-pool wiring into _init_post_memory (#31601) --- python/sglang/srt/model_executor/model_runner.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 568a83a92..39dfeda8b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -745,6 +745,11 @@ class ModelRunner: # Keep a reference so the shared byte buffer is not GC'd. self._unified_memory_pool = result.unified_memory_pool + self._init_post_memory_pool_components() + + def _init_post_memory_pool_components(self): + """Post-pool component wiring, split out of alloc_memory_pool so forks + that build bespoke memory pools can reuse it after allocating them.""" # Must be called AFTER init_memory_pool so the pool object exists for # canary to monkey-patch, and BEFORE init_decode_cuda_graph so warmup # forwards captured into the graph see the patched pool methods.