From 413913763fbeed66fcb50f184a016b1d154696d0 Mon Sep 17 00:00:00 2001 From: Fergus Date: Wed, 8 Apr 2026 10:03:39 +0100 Subject: [PATCH] fix: wrap _import_static_state in inference_mode to fix resume on Blackwell (#21035) --- .../sglang/srt/managers/scheduler_update_weights_mixin.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index abcda6794..309d3dd3c 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -220,6 +220,7 @@ def _export_static_state(model): def _import_static_state(model, static_params): - self_named_buffers = dict(model.named_buffers()) - for name, tensor in static_params["buffers"]: - self_named_buffers[name][...] = tensor + with torch.inference_mode(): + self_named_buffers = dict(model.named_buffers()) + for name, tensor in static_params["buffers"]: + self_named_buffers[name][...] = tensor