Speed up update_weights_from_tensor (#2695)

2025-01-02 18:05:19 +08:00
parent 148254d4db
commit 9183c23eca
6 changed files with 48 additions and 25 deletions
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -426,8 +426,7 @@ class UpdateWeightsFromDistributedReqOutput:

@dataclass
 class UpdateWeightsFromTensorReqInput:
-    name: str
-    tensor: torch.Tensor
+    serialized_named_tensors: bytes  # indeed Dict[str, torch.Tensor]


@dataclass