Add update_weights_from_tensor (#2631)

test/srt/run_suite.py

@@ -40,6 +40,7 @@ suites = {
         "test_triton_attention_kernels.py",
         "test_triton_attention_backend.py",
         "test_update_weights_from_disk.py",
+        "test_update_weights_from_tensor.py",
         "test_vision_chunked_prefill.py",
         "test_vision_openai_server.py",
         "test_session_control.py",

test/srt/test_update_weights_from_tensor.py

@@ -0,0 +1,32 @@
+import unittest
+
+import torch
+
+import sglang as sgl
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+
+class TestReleaseGPUOccupation(unittest.TestCase):
+    def test_release_and_resume_occupation(self):
+        engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+
+        param_name = "model.layers.2.self_attn.k_proj.weight"
+
+        def _check_param(expect_values):
+            actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5]
+            assert torch.allclose(
+                actual_values, torch.tensor(expect_values), atol=0.001
+            ), f"{actual_values=}"
+
+        _check_param([0.0571, -0.0114, 0.0444, 0.0215, -0.0149])
+
+        new_tensor = torch.full((3072, 2048), 1.5)
+        engine.update_weights_from_tensor(param_name, new_tensor)
+
+        _check_param([1.5] * 5)
+
+        engine.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main()