minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Copy decode result on forward_stream instead of copy_stream (#29642)

Browse Source
This commit is contained in:
Thomas Wang
2026-06-29 23:53:48 -07:00
committed by GitHub
parent b6fceaa789
commit f920a37da4
1 changed files with 15 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+15 -5
View File
@@ -253,6 +253,7 @@ from sglang.srt.utils import (
get_bool_env_var, get_bool_env_var,
get_int_env_var, get_int_env_var,
is_cuda, is_cuda,
is_hip,
is_mps, is_mps,
kill_itself_when_parent_died, kill_itself_when_parent_died,
require_mlp_sync, require_mlp_sync,
@@ -291,6 +292,7 @@ TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get() TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
_is_npu = is_npu() _is_npu = is_npu()
_is_hip = is_hip()
class Scheduler( class Scheduler(
@@ -3239,15 +3241,23 @@ class Scheduler(
batch_result.copy_done = self.device_module.Event() batch_result.copy_done = self.device_module.Event()
if batch_result.delay_sample_func is None: if batch_result.delay_sample_func is None:
self._relay_forward_payload(future_indices, batch_result) self._relay_forward_payload(future_indices, batch_result)
# Result D2H on copy_stream overlaps the next forward if _is_hip:
# instead of serializing on forward_stream; it's a leaf # Cross-stream sync costs more than the tiny D2H it
# gated by copy_done, so nothing on forward_stream waits. # overlaps.
self.copy_stream.wait_stream(self.forward_stream)
with self.copy_stream_ctx:
batch_result.copy_to_cpu( batch_result.copy_to_cpu(
return_logprob=batch.return_logprob, return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states, return_hidden_states=batch.return_hidden_states,
) )
else:
# Result D2H on copy_stream overlaps the next forward
# instead of serializing on forward_stream; it's a leaf
# gated by copy_done, so nothing on forward_stream waits.
self.copy_stream.wait_stream(self.forward_stream)
with self.copy_stream_ctx:
batch_result.copy_to_cpu(
return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states,
)
else: else:
batch_result.future_indices = future_indices batch_result.future_indices = future_indices