[core] Make overlap-schedule WAR barrier CUDA-only (#26646)
This commit is contained in:
@@ -106,8 +106,10 @@ class FutureMap:
|
||||
(self.req_pool_size,), dtype=torch.int64, device=self.device
|
||||
)
|
||||
# Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare
|
||||
# D2H pulls (gated only on publish, off the schedule stream).
|
||||
if _is_cuda or _is_hip:
|
||||
# D2H pulls (gated only on publish, off the schedule stream). CUDA-only:
|
||||
# recovers occupancy lost to the WAR barrier (also CUDA-only); other
|
||||
# platforms have no barrier and use the plain .cpu() bootstrap path.
|
||||
if _is_cuda:
|
||||
self.new_seq_lens_cpu_pinned = torch.empty(
|
||||
(self.req_pool_size,), dtype=torch.int64, pin_memory=True
|
||||
)
|
||||
|
||||
@@ -245,6 +245,7 @@ from sglang.srt.utils import (
|
||||
get_available_gpu_memory,
|
||||
get_bool_env_var,
|
||||
get_int_env_var,
|
||||
is_cuda,
|
||||
is_mps,
|
||||
kill_itself_when_parent_died,
|
||||
require_mlp_sync,
|
||||
@@ -1362,6 +1363,8 @@ class Scheduler(
|
||||
self.schedule_stream = self.device_module.Stream(priority=0)
|
||||
if self.device == "cpu":
|
||||
self.schedule_stream.synchronize = lambda: None # No-op for CPU
|
||||
# WAR barrier is CUDA-only; other platforms keep the pre-barrier behavior.
|
||||
self._war_barrier_enabled = is_cuda()
|
||||
with self.device_module.StreamContext(self.schedule_stream):
|
||||
dispatch_event_loop(self)
|
||||
|
||||
@@ -1412,7 +1415,8 @@ class Scheduler(
|
||||
continue
|
||||
|
||||
# WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads.
|
||||
self.schedule_stream.wait_stream(self.forward_stream)
|
||||
if self._war_barrier_enabled:
|
||||
self.schedule_stream.wait_stream(self.forward_stream)
|
||||
|
||||
# Get the next batch to run
|
||||
batch = self.get_next_batch_to_run()
|
||||
|
||||
Reference in New Issue
Block a user