[Feature] Distributed layerwise offload with DP multi-concurrency + mmap weight… (#5397)

Signed-off-by: evanchueng <evanchueng@outlook.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: david6666666 <530634352@qq.com>
This commit is contained in:
evanchuenganddavid6666666 authored and GitHub committed 2026-07-30 13:35:38 +00:00
1 parent b3f4fbf9e2
commit 744c65b730
23 files changed
+3586 -78

No files matched your search

@@ -102,7 +102,10 @@ def _start_worker(req_q, res_q, count=2):
def _run():
for _ in range(count):
req = req_q.get(timeout=10)
try:
req = req_q.get(timeout=10)
except queue.Empty:
break
method = req.get("method", "")
args = req.get("args", ())
if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput):
@@ -146,7 +149,7 @@ def _inject_interleave(executor):
orig_enqueue(item)
if threading.current_thread().name == "thread_a":
a_enqueued.set() # tell B: "A has enqueued"
b_complete.wait(5) # block A until B finishes
b_complete.wait(1) # block A until B finishes
executor._broadcast_mq.enqueue = _controlled
return a_enqueued, b_complete
@@ -173,8 +176,8 @@ class TestConcurrentRequestExecution:
results["B"] = engine.add_req_and_wait_for_response(_mock_request("B"))
b_complete.set() # release A
ta = threading.Thread(target=_a, name="thread_a")
tb = threading.Thread(target=_b, name="thread_b")
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
ta.start()
tb.start()
ta.join(10)
@@ -274,8 +277,8 @@ class TestConcurrentCollectiveRpc:
)
b_complete.set()
ta = threading.Thread(target=_a, name="thread_a")
tb = threading.Thread(target=_b, name="thread_b")
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
ta.start()
tb.start()
ta.join(10)
@@ -311,8 +314,8 @@ class TestConcurrentRequestExecutionAndCollectiveRpc:
)
b_complete.set()
ta = threading.Thread(target=_a, name="thread_a")
tb = threading.Thread(target=_b, name="thread_b")
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
ta.start()
tb.start()
ta.join(10)
@@ -683,7 +686,10 @@ class TestMultiprocExecutorStepStreamingOutput:
def _worker():
for runner_output in runner_outputs:
req_q.get(timeout=10)
try:
req_q.get(timeout=10)
except queue.Empty:
break
res_q.put(runner_output)
thread = threading.Thread(target=_worker, daemon=True)