[feat] support in-flight weight update (#10071)

Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>

python/sglang/srt/entrypoints/http_server.py

@@ -78,6 +78,7 @@ from sglang.srt.managers.io_struct import (
     AbortReq,
     CloseSessionReqInput,
     ConfigureLoggingReq,
+    ContinueGenerationReqInput,
     DestroyWeightsUpdateGroupReqInput,
     EmbeddingReqInput,
     GenerateReqInput,
@@ -87,6 +88,7 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqInput,
     OpenSessionReqInput,
     ParseFunctionCallReq,
+    PauseGenerationReqInput,
     ProfileReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
@@ -1087,9 +1089,9 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re
 
 
 @app.post("/pause_generation")
-async def pause_generation(request: Request):
+async def pause_generation(obj: PauseGenerationReqInput, request: Request):
     """Pause generation."""
-    await _global_state.tokenizer_manager.pause_generation()
+    await _global_state.tokenizer_manager.pause_generation(obj)
     return ORJSONResponse(
         content={"message": "Generation paused successfully.", "status": "ok"},
         status_code=200,
@@ -1097,9 +1099,9 @@ async def pause_generation(request: Request):
 
 
 @app.post("/continue_generation")
-async def continue_generation(request: Request):
+async def continue_generation(obj: ContinueGenerationReqInput, request: Request):
     """Continue generation."""
-    await _global_state.tokenizer_manager.continue_generation()
+    await _global_state.tokenizer_manager.continue_generation(obj)
     return ORJSONResponse(
         content={"message": "Generation continued successfully.", "status": "ok"},
         status_code=200,

python/sglang/srt/managers/io_struct.py

@@ -21,7 +21,7 @@ import uuid
 from abc import ABC
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.schedule_batch import BaseFinishReason
@@ -1064,6 +1064,41 @@ class FlushCacheReqOutput(BaseReq):
     success: bool
 
 
+@dataclass
+class PauseGenerationReqInput(BaseReq):
+    """
+    Note that the PauseGenerationRequests is only supported in SGLang Server.
+    abort: Abort and return all requests currently being processed.
+
+    in_place: Pause the scheduler's event_loop from performing inference;
+            only non-inference requests (e.g., control commands) will be handled.
+            The requests in the engine will be paused and stay in the event_loop,
+            then continue generation after continue_generation with the old kv cache.
+            Note: In 'inplace' mode, flush_cache will fail if there are any requests
+            in the running_batch.
+
+    retract: Pause the scheduler's event loop from performing inference;
+            only non-inference requests will be handled, and all currently running
+            requests will be retracted back to the waiting_queue.
+            Note: The KV cache can be flushed in this mode and will be automatically
+            recomputed after continue_generation.
+    """
+
+    mode: Literal["abort", "retract", "in_place"] = "abort"
+
+    def __post_init__(self):
+        allowed = ["abort", "retract", "in_place"]
+        if self.mode not in allowed:
+            raise ValueError(
+                f"Invalid mode: {self.mode!r}. " f"Expected one of {allowed}."
+            )
+
+
+@dataclass
+class ContinueGenerationReqInput(BaseReq):
+    pass
+
+
 @dataclass
 class UpdateWeightFromDiskReqInput(BaseReq):
     # The model path with the new weights
@@ -1084,6 +1119,8 @@ class UpdateWeightFromDiskReqInput(BaseReq):
     recapture_cuda_graph: bool = False
     # The trainer step id. Used to know which step's weights are used for sampling.
     token_step: int = 0
+    # Whether to flush the cache after updating weights
+    flush_cache: bool = True
 
 
 @dataclass

python/sglang/srt/managers/schedule_batch.py

@@ -1512,8 +1512,17 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         evict_from_tree_cache(self.tree_cache, num_tokens)
         return self._is_available_size_sufficient(num_tokens)
 
+    def retract_all(self, server_args: ServerArgs):
+        retracted_reqs = self.reqs
+        for idx in range(len(self.reqs)):
+            self.release_req(idx, len(self.reqs) - idx, server_args)
+
+        self.filter_batch(retracted_reqs)
+        return retracted_reqs
+
     def retract_decode(
-        self, server_args: ServerArgs
+        self,
+        server_args: ServerArgs,
     ) -> Tuple[List[Req], float, List[Req]]:
         """Retract the decoding requests when there is not enough memory."""
         sorted_indices = list(range(len(self.reqs)))

python/sglang/srt/managers/scheduler.py

@@ -73,6 +73,7 @@ from sglang.srt.managers.io_struct import (
     ClearHiCacheReqInput,
     ClearHiCacheReqOutput,
     CloseSessionReqInput,
+    ContinueGenerationReqInput,
     DestroyWeightsUpdateGroupReqInput,
     ExpertDistributionReq,
     ExpertDistributionReqOutput,
@@ -93,6 +94,7 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterReqOutput,
     OpenSessionReqInput,
     OpenSessionReqOutput,
+    PauseGenerationReqInput,
     ProfileReq,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
@@ -443,6 +445,7 @@ class Scheduler(
         if self.device == "cpu":
             self.default_stream.synchronize = lambda: None  # No-op for CPU
         self.forward_sleep_time = None
+        self._engine_paused = False
 
         # Init chunked prefill
         self.chunked_prefill_size = server_args.chunked_prefill_size
@@ -568,6 +571,8 @@ class Scheduler(
                 (LoadLoRAAdapterReqInput, self.load_lora_adapter),
                 (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
                 (GetLoadReqInput, self.get_load),
+                (PauseGenerationReqInput, self.pause_generation),
+                (ContinueGenerationReqInput, self.continue_generation),
             ]
         )
 
@@ -953,6 +958,9 @@ class Scheduler(
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
 
+            if self._engine_paused:
+                continue
+
             batch = self.get_next_batch_to_run()
             self.cur_batch = batch
 
@@ -985,6 +993,9 @@ class Scheduler(
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
 
+            if self._engine_paused:
+                continue
+
             batch = self.get_next_batch_to_run()
             self.cur_batch = batch
 
@@ -2154,8 +2165,7 @@ class Scheduler(
 
     def _is_no_request(self):
         no_request = (
-            len(self.waiting_queue) == 0
+            self.running_batch.is_empty()
-            and self.running_batch.is_empty()
             and (self.last_batch is None or self.last_batch.is_empty())
             and (self.cur_batch is None or self.cur_batch.is_empty())
             and (not self.enable_overlap or len(self.result_queue) == 0)
@@ -2428,6 +2438,29 @@ class Scheduler(
     def _pause_engine(self) -> Tuple[List[Req], int]:
         raise NotImplementedError()
 
+    def pause_generation(self, recv_req: PauseGenerationReqInput):
+        self._engine_paused = True
+
+        if self.enable_overlap and self.last_batch:
+            # Process the results of the last batch
+            tmp_batch, tmp_result = self.result_queue.popleft()
+            self.process_batch_result(tmp_batch, tmp_result)
+            self.last_batch = None
+            self.cur_batch = None
+
+        if recv_req.mode == "retract":
+            self.running_batch.filter_batch()
+            if len(self.running_batch.reqs) != 0:
+                retracted_reqs = self.running_batch.retract_all(self.server_args)
+                for req in retracted_reqs:
+                    self._add_request_to_queue(req)
+
+            self.running_batch.batch_is_full = False
+            self.chunked_req = None
+
+    def continue_generation(self, recv_req: ContinueGenerationReqInput):
+        self._engine_paused = False
+
     def load_lora_adapter(
         self, recv_req: LoadLoRAAdapterReqInput
     ) -> LoadLoRAAdapterReqOutput:

python/sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -44,6 +44,7 @@ class SchedulerUpdateWeightsMixin:
         """In-place update of the weights from disk."""
         success, message = self.tp_worker.update_weights_from_disk(recv_req)
         if success:
+            if recv_req.flush_cache:
                 flush_cache_success = self.flush_cache()
                 assert flush_cache_success, "Cache flush failed after updating weights"
         else:

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -404,6 +404,14 @@ class TokenizerCommunicatorMixin:
         if obj.abort_all_requests:
             self.abort_request(abort_all=True)
 
+        # Immediately update the weights if the engine is in paused state
+        async with self.is_pause_cond:
+            if self.is_pause:
+                result = (await self.update_weights_from_distributed_communicator(obj))[
+                    0
+                ]
+                return result.success, result.message
+
         # This means that weight sync
         # cannot run while requests are in progress.
         async with self.model_update_lock.writer_lock:
@@ -457,6 +465,12 @@ class TokenizerCommunicatorMixin:
         if obj.abort_all_requests:
             self.abort_request(abort_all=True)
 
+        # Immediately update the weights if the engine is in paused state
+        async with self.is_pause_cond:
+            if self.is_pause:
+                result = (await self.update_weights_from_tensor_communicator(obj))[0]
+                return result.success, result.message
+
         # This means that weight sync
         # cannot run while requests are in progress.
         async with self.model_update_lock.writer_lock:

python/sglang/srt/managers/tokenizer_manager.py

@@ -55,6 +55,7 @@ from sglang.srt.managers.io_struct import (
     BatchTokenizedEmbeddingReqInput,
     BatchTokenizedGenerateReqInput,
     ConfigureLoggingReq,
+    ContinueGenerationReqInput,
     EmbeddingReqInput,
     FreezeGCReq,
     GenerateReqInput,
@@ -62,6 +63,7 @@ from sglang.srt.managers.io_struct import (
     HealthCheckOutput,
     LoadLoRAAdapterReqInput,
     OpenSessionReqOutput,
+    PauseGenerationReqInput,
     SessionParams,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -1246,9 +1248,12 @@ class TokenizerManager(TokenizerCommunicatorMixin):
                 self.metrics_collector.labels
             )
 
-    async def pause_generation(self):
+    async def pause_generation(self, obj: PauseGenerationReqInput):
         async with self.is_pause_cond:
             self.is_pause = True
+            if obj.mode != "abort":
+                await self.send_to_scheduler.send_pyobj(obj)
+            else:
                 # we are using the model_update_lock to check if there is still on-going requests.
                 while True:
                     # TODO: maybe make it async instead of fire-and-forget
@@ -1258,9 +1263,10 @@ class TokenizerManager(TokenizerCommunicatorMixin):
                         break
                     await asyncio.sleep(1.0)
 
-    async def continue_generation(self):
+    async def continue_generation(self, obj: ContinueGenerationReqInput):
         async with self.is_pause_cond:
             self.is_pause = False
+            await self.send_to_scheduler.send_pyobj(obj)
             self.is_pause_cond.notify_all()
 
     async def update_weights_from_disk(
@@ -1278,6 +1284,11 @@ class TokenizerManager(TokenizerCommunicatorMixin):
         if obj.abort_all_requests:
             self.abort_request(abort_all=True)
 
+        # Immediately update the weights if the engine is in paused state
+        async with self.is_pause_cond:
+            if self.is_pause:
+                return await self._wait_for_model_update_from_disk(obj)
+
         if True:  # Keep this redundant check to simplify some internal code sync
             # Hold the lock if it is not async. This means that weight sync
             # cannot run while requests are in progress.

test/srt/rl/test_update_weights_from_disk.py

@@ -98,19 +98,51 @@ class TestServerUpdateWeightsFromDisk(CustomTestCase):
         print(f"[Server Mode] Generated text: {response.json()['text']}")
         return response.json()["text"]
 
+    def run_decode_random(self, max_new_tokens=32):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": f"Question: {random.randint(0, 100)},The capital of France is",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": max_new_tokens,
+                    "ignore_eos": True,
+                },
+            },
+        )
+        return response.json()
+
     def get_model_info(self):
         response = requests.get(self.base_url + "/get_model_info")
         model_path = response.json()["model_path"]
         print(json.dumps(response.json()))
         return model_path
 
-    def run_update_weights(self, model_path):
+    def run_update_weights(self, model_path, flush_cache=True):
         response = requests.post(
             self.base_url + "/update_weights_from_disk",
-            json={"model_path": model_path},
+            json={
+                "model_path": model_path,
+                "flush_cache": flush_cache,
+            },
+        )
+        ret = response.json()
+        return ret
+
+    def pause_generation(self, mode):
+        response = requests.post(
+            self.base_url + "/pause_generation",
+            json={"mode": mode},
+        )
+        ret = response.json()
+        return ret
+
+    def continue_generation(self):
+        response = requests.post(
+            self.base_url + "/continue_generation",
+            json={},
         )
         ret = response.json()
-        print(json.dumps(ret))
         return ret
 
     def test_update_weights(self):
@@ -138,6 +170,42 @@ class TestServerUpdateWeightsFromDisk(CustomTestCase):
         updated_response = self.run_decode()
         self.assertEqual(origin_response[:32], updated_response[:32])
 
+    def test_update_weights_non_blocking(self):
+        origin_model_path = self.get_model_info()
+        print(f"[Server Mode] origin_model_path: {origin_model_path}")
+
+        pause_generation_modes = ["in_place", "retract"]
+        for pause_generation_mode in pause_generation_modes:
+            num_requests = 32
+            with ThreadPoolExecutor(num_requests) as executor:
+                futures = [
+                    executor.submit(self.run_decode_random, 1600)
+                    for _ in range(num_requests)
+                ]
+
+                # ensure the decode has been started
+                time.sleep(2)
+
+                new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace(
+                    "-Instruct", ""
+                )
+                ret = self.pause_generation(pause_generation_mode)
+                ret = self.run_update_weights(
+                    new_model_path, flush_cache=pause_generation_mode == "retract"
+                )
+                self.assertTrue(ret["success"])
+                ret = self.continue_generation()
+
+                for future in as_completed(futures):
+                    self.assertNotEqual(
+                        future.result()["meta_info"]["finish_reason"]["type"], "abort"
+                    )
+
+            updated_model_path = self.get_model_info()
+            print(f"[Server Mode] updated_model_path: {updated_model_path}")
+            self.assertEqual(updated_model_path, new_model_path)
+            self.assertNotEqual(updated_model_path, origin_model_path)
+
     def test_update_weights_unexist_model(self):
         origin_model_path = self.get_model_info()
         print(f"[Server Mode] origin_model_path: {origin_model_path}")

test/srt/rl/test_update_weights_from_distributed.py

@@ -18,6 +18,7 @@ import os
 import random
 import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 import requests
@@ -68,6 +69,8 @@ def init_process(
     backend,
     checking_parameters,
     tie_word_embeddings,
+    barrier,
+    pause_generation_mode,
 ):
     torch.cuda.set_device(rank)
 
@@ -81,6 +84,7 @@ def init_process(
             checking_parameters,
             tie_word_embeddings,
             state_dict_key_to_shape,
+            barrier,
         )
     elif rank in [1, 2]:
         init_process_sgl(
@@ -94,6 +98,8 @@ def init_process(
             state_dict_key_to_shape,
             backend,
             tp_size,
+            barrier,
+            pause_generation_mode,
         )
 
 
@@ -106,6 +112,7 @@ def init_process_hf(
     checking_parameters,
     tie_word_embeddings,
     state_dict_key_to_shape,
+    barrier,
 ):
     # These two environment variables are very important
     # to avoid unexpected behaviors of CUDA and NCCL.
@@ -162,6 +169,7 @@ def init_process_hf(
         group_name="test_parameter_update_group",
     )
     torch.cuda.synchronize()
+    barrier.wait()
     time_begin_broadcast = time.perf_counter()
 
     # The last parameter is lm_head.weight, which is tied
@@ -208,6 +216,8 @@ def init_process_sgl(
     state_dict_key_to_shape,
     backend,
     tp_size,
+    barrier,
+    pause_generation_mode,
 ):
     torch.cuda.set_device(rank)
     torch.cuda.synchronize()
@@ -282,8 +292,25 @@ def init_process_sgl(
             },
         )
 
-    torch.cuda.synchronize()
+    if pause_generation_mode in ["in_place", "retract"]:
-    time_begin_update = time.perf_counter()
+
+        def run_decode(max_new_tokens=32):
+            response = requests.post(
+                url + "/generate",
+                json={
+                    "text": f"Question: {random.randint(0, 100)},The capital of France is",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                },
+            )
+            return response.json()
+
+        with ThreadPoolExecutor(32) as executor:
+            futures = [executor.submit(run_decode, 1000) for _ in range(32)]
+            time.sleep(2)
 
     # The last parameter is lm_head.weight, which is tied
     # with embed_tokens.weight. Actually, we only need
@@ -300,6 +327,14 @@ def init_process_sgl(
     dtypes = [torch.bfloat16 if backend == "Engine" else "bfloat16"] * len(names)
     shapes = [state_dict_key_to_shape[parameter_name] for parameter_name in names]
 
+    if pause_generation_mode in ["in_place", "retract"]:
+        requests.post(
+            url + "/pause_generation",
+            json={"mode": pause_generation_mode},
+        )
+    torch.cuda.synchronize()
+    barrier.wait()
+    time_begin_update = time.perf_counter()
     if backend == "Engine":
         engine.update_weights_from_distributed(
             names,
@@ -315,10 +350,23 @@ def init_process_sgl(
                 "dtypes": dtypes,
                 "shapes": shapes,
                 "group_name": "test_parameter_update_group",
+                "flush_cache": not (pause_generation_mode == "in_place"),
             },
         )
     torch.cuda.synchronize()
     time_end_update = time.perf_counter()
+    if pause_generation_mode in ["in_place", "retract"]:
+        requests.post(
+            url + "/continue_generation",
+            json={},
+        )
+
+        # discard unfinished requests to save test overhead
+        time.sleep(2)
+        requests.post(
+            url + "/pause_generation",
+            json={"mode": "abort"},
+        )
 
     # Measure the latency of broadcast/weights update.
     update_time = time_end_update - time_begin_update
@@ -383,6 +431,7 @@ def test_update_weights_from_distributed(
     state_dict_key_to_shape,
     truncate_size,
     checking_parameters,
+    pause_generation_mode=None,
 ):
     tie_word_embeddings = (
         True if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else False
@@ -393,6 +442,7 @@ def test_update_weights_from_distributed(
     )
     param_queue = mp.Queue()
     results = {}
+    barrier = mp.Barrier(1 + dp_size)
 
     context = mp.spawn(
         init_process,
@@ -406,6 +456,8 @@ def test_update_weights_from_distributed(
             backend,
             checking_parameters,
             tie_word_embeddings,
+            barrier,
+            pause_generation_mode,
         ),
         nprocs=1 + dp_size,
         join=False,
@@ -558,28 +610,50 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
         # test_suits : tp, dp, model_name, backend
         if is_in_ci():
             mode = random.choice(["Engine", "Server"])
+            if mode == "Server":
+                pause_generation_mode = random.choice(["in_place", "retract"])
+            else:
+                pause_generation_mode = None
             test_suits = [
-                (1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode),
+                (1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode, pause_generation_mode),
             ]
         else:
             test_suits = [
-                (1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
+                (1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
-                (1, 1, DEFAULT_MODEL_NAME_FOR_TEST, "Sever"),
+                (
+                    1,
+                    1,
+                    DEFAULT_MODEL_NAME_FOR_TEST,
+                    "Sever",
+                    random.choice(["in_place", "retract"]),
+                ),
             ]
 
             if torch.cuda.device_count() >= 4:
                 test_suits.extend(
                     [
-                        (2, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
+                        (2, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
-                        (1, 2, DEFAULT_MODEL_NAME_FOR_TEST, "Server"),
+                        (
+                            1,
+                            2,
+                            DEFAULT_MODEL_NAME_FOR_TEST,
+                            "Server",
+                            random.choice(["in_place", "retract"]),
+                        ),
                     ]
                 )
 
             if torch.cuda.device_count() >= 5:
                 test_suits.extend(
                     [
-                        (2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
+                        (2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
-                        (2, 2, DEFAULT_MODEL_NAME_FOR_TEST, "Server"),
+                        (
+                            2,
+                            2,
+                            DEFAULT_MODEL_NAME_FOR_TEST,
+                            "Server",
+                            random.choice(["in_place", "retract"]),
+                        ),
                     ]
                 )
 
@@ -615,7 +689,7 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
             "lm_head.weight",
         ]
 
-        for tp_size, dp_size, model_name, backend in test_suits:
+        for tp_size, dp_size, model_name, backend, pause_generation_mode in test_suits:
             test_update_weights_from_distributed(
                 tp_size,
                 dp_size,
@@ -624,6 +698,7 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
                 model_state_dict_shapes[model_name],
                 truncate_size,
                 checking_parameters,
+                pause_generation_mode,
             )
 
 

test/srt/rl/test_update_weights_from_tensor.py

@@ -1,12 +1,23 @@
 import gc
+import json
+import random
 import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
+import requests
 import torch
 
 import sglang as sgl
+from sglang.srt.utils import MultiprocessingSerializer, kill_process_tree
 from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
-from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
 
 
 def test_update_weights_from_tensor(tp_size):
@@ -167,6 +178,112 @@ class TestUpdateWeightsFromTensor(CustomTestCase):
         engine.shutdown()
 
 
+class TestServerUpdateWeightsFromTensorNonBlocking(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--max-running-requests", 8],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_decode(self, max_new_tokens=32):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": f"Question: {random.randint(0, 100)},The capital of France is",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": max_new_tokens,
+                    "ignore_eos": True,
+                },
+            },
+        )
+        return response.json()
+
+    def get_model_info(self):
+        response = requests.get(self.base_url + "/get_model_info")
+        model_path = response.json()["model_path"]
+        print(json.dumps(response.json()))
+        return model_path
+
+    def pause_generation(self, mode):
+        response = requests.post(
+            self.base_url + "/pause_generation",
+            json={"mode": mode},
+        )
+        ret = response.json()
+        return ret
+
+    def continue_generation(self):
+        response = requests.post(
+            self.base_url + "/continue_generation",
+            json={},
+        )
+        ret = response.json()
+        return ret
+
+    def run_update_weights(self, named_tensors, flush_cache=True):
+        response = requests.post(
+            self.base_url + "/update_weights_from_tensor",
+            json={
+                "serialized_named_tensors": [
+                    MultiprocessingSerializer.serialize(named_tensors, output_str=True)
+                ],
+                "flush_cache": flush_cache,
+            },
+        )
+        ret = response.json()
+        return ret
+
+    def test_update_weights(self):
+        pause_generation_modes = ["in_place", "retract"]
+        for pause_generation_mode in pause_generation_modes:
+            num_requests = 32
+            with ThreadPoolExecutor(num_requests) as executor:
+                futures = [
+                    executor.submit(self.run_decode, 3000) for _ in range(num_requests)
+                ]
+
+                # ensure the decode has been started
+                time.sleep(2)
+
+                param_names = [
+                    f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)
+                ]
+                new_tensor = torch.full((16384, 2048), 1.5, device="cuda")
+                named_tensors = [(x, new_tensor) for x in param_names]
+
+                ret = self.pause_generation(pause_generation_mode)
+                ret = self.run_update_weights(
+                    named_tensors, flush_cache=pause_generation_mode == "retract"
+                )
+                self.assertTrue(ret["success"])
+                ret = self.continue_generation()
+
+                for future in as_completed(futures):
+                    self.assertNotEqual(
+                        future.result()["meta_info"]["finish_reason"]["type"], "abort"
+                    )
+
+                for param_name in param_names[:3]:
+                    response = requests.post(
+                        self.base_url + "/get_weights_by_name",
+                        json={"name": param_name},
+                    )
+                    actual_values = torch.tensor(response.json())[0, :5]
+                    assert torch.allclose(
+                        actual_values, torch.tensor([1.5] * 5), atol=0.002
+                    ), f"{actual_values=}"
+
+
 def _check_param(engine, param_name, expect_values):
     actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5]
     assert torch.allclose(