Fix async loading of RunAI-streamed tensors (#32896)

Co-authored-by: Danila Shtan <dan@nebius.com>
This commit is contained in:
Danila Shtan
2026-07-31 21:46:33 +08:00
committed by GitHub
co-authored by Danila Shtan
parent 690de097c4
commit 5f9b0db18c
2 changed files with 44 additions and 0 deletions
+6
View File
@@ -285,7 +285,13 @@ def should_async_load(weight: torch.Tensor) -> bool:
For host (CPU) tensors, using a threadpool can overlap H2D copies
and improve throughput. For device tensors, threading often adds overhead
(e.g., GIL contention) without benefit, so we do it synchronously.
RunAI-streamed tensors are zero-copy views into a reused CPU buffer. They
must be consumed synchronously before the streamer fills its next batch.
"""
if getattr(weight, "_sglang_runai_streamer_tensor", False):
return False
device = getattr(weight, "device", None)
if device is None:
return False