[Bugfix] Export weight-cache IPC tensors separately for each client (#56472)
Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
shaohuaxi committed
Sep 14, 2026 at 19:43 UTC
567f745c1b14ffaf09bfa71b7c4d0177a7994c04
1 file changed
+13
-18
vllm/model_executor/model_loader/weight_cache/daemon.py
+13
-18
@@ -93,11 +93,14 @@ def export_entries(
93
``named_parameters``/``named_buffers`` are iterated with
94
``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing
95
storage with ``embed_tokens.weight``) are not silently dropped. Each unique
96
- tensor is exported once; every additional name that refers to the same
96
+ tensor is exported once per call; every additional name that refers to the same
97
tensor object is recorded in the returned alias map so the client can
98
re-establish the shared identity instead of allocating uninitialized
99
memory for it.
100
101
+ CUDA reduction arguments must be exported separately for each consumer so
102
+ that PyTorch registers a reference for each IPC mapping's lifetime.
103
+
104
Returns:
105
A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to
106
its ``TensorEntry`` and ``aliases`` maps each duplicate name to its
@@ -167,8 +170,6 @@ class WeightCacheDaemon:
170
self.distributed_init_method = distributed_init_method
171
self.socket_dir = socket_dir
172
self.model: torch.nn.Module | None = None
170
- self.entries: dict[str, TensorEntry] = {}
171
- self.aliases: dict[str, str] = {}
173
# Fingerprint before loading: process_weights_after_loading may
174
# mutate hf_config.quantization_config.
175
self.cache_config = WeightCacheKey.from_model_config(
@@ -190,17 +191,11 @@ class WeightCacheDaemon:
191
with set_current_vllm_config(self.vllm_config):
192
ensure_model_parallel_initialized(tp_size, 1)
193
self.model = get_daemon_model(self.vllm_config)
193
- self._export_entries()
194
logger.info(
195
- "Weight cache daemon rank %d cached %d tensors",
195
+ "Weight cache daemon rank %d loaded model",
196
self.tp_rank,
197
- len(self.entries),
197
)
198
200
- def _export_entries(self) -> None:
201
- assert self.model is not None
202
- self.entries, self.aliases = export_entries(self.model)
203
-
199
def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None:
200
"""Serve requests until terminated.
201
@@ -295,28 +290,28 @@ class WeightCacheDaemon:
290
logger.warning("WeightCacheKey mismatch on fields: %s", mismatched)
291
send_msg(conn, {"status": "mismatch", "fields": mismatched})
292
return
298
- if not self.entries:
293
+ if self.model is None:
294
send_msg(conn, {"status": "error", "message": "Weights were released"})
295
return
296
+ gpu_uuid = get_current_device_uuid()
297
+ entries, aliases = export_entries(self.model)
298
send_msg(
299
conn,
300
{
301
"status": "ok",
305
- "entries": self.entries,
306
- "aliases": self.aliases,
307
- "gpu_uuid": get_current_device_uuid(),
302
+ "entries": entries,
303
+ "aliases": aliases,
304
+ "gpu_uuid": gpu_uuid,
305
},
306
)
307
logger.info_once(
308
"Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine",
309
self.tp_rank,
313
- len(self.entries),
314
- len(self.aliases),
310
+ len(entries),
311
+ len(aliases),
312
)
313
314
def _handle_release(self, conn: socket.socket) -> None:
318
- self.entries.clear()
319
- self.aliases.clear()
315
self.model = None
316
torch.accelerator.empty_cache()
317
logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank)