diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index fad593f62..1611b8fcc 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -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 diff --git a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py index d4fc161b8..800fdc101 100644 --- a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py +++ b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py @@ -1,4 +1,6 @@ +import concurrent.futures import sys +import threading import unittest from types import SimpleNamespace from typing import cast @@ -7,6 +9,7 @@ from unittest.mock import patch import torch import sglang.srt.model_loader.loader as loader_mod +import sglang.srt.model_loader.utils as model_loader_utils import sglang.srt.model_loader.weight_utils as weight_utils from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig, LoadFormat @@ -110,6 +113,41 @@ class TestRunaiModelStreamerLoader(CustomTestCase): marked.fill_(2) self.assertEqual(cloned.item(), 1) + def test_runai_streamed_tensor_is_consumed_before_buffer_reuse(self): + def consume_view(mark_as_runai: bool): + shared_buffer = torch.tensor([1], dtype=torch.int32) + view = shared_buffer[:] + if mark_as_runai: + setattr(view, weight_utils.RUNAI_STREAMER_TENSOR_ATTR, True) + + release_worker = threading.Event() + observed = [] + futures = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + # Keep the sole worker busy so an async consumer cannot read the + # zero-copy view until after the simulated streamer buffer reuse. + blocker = executor.submit(release_worker.wait) + model_loader_utils.maybe_executor_submit( + executor=executor, + futures=futures, + use_async=model_loader_utils.should_async_load(view), + func=lambda tensor: observed.append(tensor.item()), + func_args=(view,), + ) + shared_buffer.fill_(2) + release_worker.set() + blocker.result() + for future in futures: + future.result() + + return observed, len(futures) + + # The control demonstrates the race: an async consumer observes the + # overwritten buffer rather than the value present when it was queued. + self.assertEqual(consume_view(mark_as_runai=False), ([2], 1)) + # A RunAI-tagged view is consumed inline before the buffer is reused. + self.assertEqual(consume_view(mark_as_runai=True), ([1], 0)) + def test_deepseek_v4_streaming_dequant_fp8_wo_a_pairs_weight_and_scale(self): weight = torch.eye(128, dtype=torch.float32).to(torch.float8_e4m3fn) scale = torch.ones((1, 1), dtype=torch.float32)