[core] Make overlap-schedule WAR barrier CUDA-only (#26646)
This commit is contained in:
@@ -1629,7 +1629,8 @@ class SchedulerDisaggregationDecodeMixin:
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
# WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads.
|
# 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
|
# Get the next batch to run
|
||||||
batch = self.get_next_disagg_decode_batch_to_run()
|
batch = self.get_next_disagg_decode_batch_to_run()
|
||||||
|
|||||||
@@ -437,7 +437,8 @@ class SchedulerDisaggregationPrefillMixin:
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
# WAR barrier on shared GPU buffers (req_to_token_pool / SWA mapping).
|
# WAR barrier on shared GPU buffers (req_to_token_pool / SWA mapping).
|
||||||
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
|
# Get the next batch to run
|
||||||
batch = self.get_next_disagg_prefill_batch_to_run()
|
batch = self.get_next_disagg_prefill_batch_to_run()
|
||||||
|
|||||||
@@ -106,8 +106,10 @@ class FutureMap:
|
|||||||
(self.req_pool_size,), dtype=torch.int64, device=self.device
|
(self.req_pool_size,), dtype=torch.int64, device=self.device
|
||||||
)
|
)
|
||||||
# Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare
|
# Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare
|
||||||
# D2H pulls (gated only on publish, off the schedule stream).
|
# D2H pulls (gated only on publish, off the schedule stream). CUDA-only:
|
||||||
if _is_cuda or _is_hip:
|
# 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.new_seq_lens_cpu_pinned = torch.empty(
|
||||||
(self.req_pool_size,), dtype=torch.int64, pin_memory=True
|
(self.req_pool_size,), dtype=torch.int64, pin_memory=True
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -245,6 +245,7 @@ from sglang.srt.utils import (
|
|||||||
get_available_gpu_memory,
|
get_available_gpu_memory,
|
||||||
get_bool_env_var,
|
get_bool_env_var,
|
||||||
get_int_env_var,
|
get_int_env_var,
|
||||||
|
is_cuda,
|
||||||
is_mps,
|
is_mps,
|
||||||
kill_itself_when_parent_died,
|
kill_itself_when_parent_died,
|
||||||
require_mlp_sync,
|
require_mlp_sync,
|
||||||
@@ -1362,6 +1363,8 @@ class Scheduler(
|
|||||||
self.schedule_stream = self.device_module.Stream(priority=0)
|
self.schedule_stream = self.device_module.Stream(priority=0)
|
||||||
if self.device == "cpu":
|
if self.device == "cpu":
|
||||||
self.schedule_stream.synchronize = lambda: None # No-op for 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):
|
with self.device_module.StreamContext(self.schedule_stream):
|
||||||
dispatch_event_loop(self)
|
dispatch_event_loop(self)
|
||||||
|
|
||||||
@@ -1412,7 +1415,8 @@ class Scheduler(
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
# WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads.
|
# 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
|
# Get the next batch to run
|
||||||
batch = self.get_next_batch_to_run()
|
batch = self.get_next_batch_to_run()
|
||||||
|
|||||||
Reference in New Issue
Block a user