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

[AMD][DSV4] Remove per-batch D2H syncs in MTP to avoid bubbles between 2 batches (#29420)

Browse Source 
Co-authored-by: amd-danli103 <amd-danli103@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
amd-danli103
2026-06-29 22:39:59 -07:00
committed by GitHub
co-authored by amd-danli103 Thomas Wang
parent f41d455f37
commit 54e71506b3
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+9 -1
View File
@@ -589,11 +589,13 @@ class DeepseekV4HipRadixBackend(
out_cache_loc: Optional[torch.Tensor] = None, out_cache_loc: Optional[torch.Tensor] = None,
extend_seq_lens: Optional[torch.Tensor] = None, extend_seq_lens: Optional[torch.Tensor] = None,
use_prefill_cuda_graph: bool = False, use_prefill_cuda_graph: bool = False,
seq_lens_cpu: Optional[List[int]] = None,
) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]:
# HIP path: build target-verify metadata eagerly even when # HIP path: build target-verify metadata eagerly even when
# SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can # SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can
# hit planner invariants during graph capture for DSV4+EAGLE. # hit planner invariants during graph capture for DSV4+EAGLE.
seq_lens_cpu = seq_lens.tolist() if seq_lens_cpu is None:
seq_lens_cpu = seq_lens.tolist()
return self.init_forward_metadata_target_verify_old( return self.init_forward_metadata_target_verify_old(
max_seq_len=max_seq_len, max_seq_len=max_seq_len,
req_pool_indices=req_pool_indices, req_pool_indices=req_pool_indices,
@@ -876,6 +878,9 @@ class DeepseekV4HipRadixBackend(
seq_lens=seq_lens, seq_lens=seq_lens,
out_cache_loc=out_cache_loc_padded, out_cache_loc=out_cache_loc_padded,
use_prefill_cuda_graph=True, use_prefill_cuda_graph=True,
# CPU mirror already available here (== seq_lens, no D2H);
# pass it so target_verify skips the per-iter seq_lens.tolist() sync.
seq_lens_cpu=seq_lens_cpu.tolist(),
) )
elif bucket == _GraphBucket.DRAFT_EXTEND: elif bucket == _GraphBucket.DRAFT_EXTEND:
num_tokens_per_bs = self.draft_extend_num_tokens_per_bs num_tokens_per_bs = self.draft_extend_num_tokens_per_bs
@@ -952,6 +957,9 @@ class DeepseekV4HipRadixBackend(
seq_lens=seq_lens, seq_lens=seq_lens,
out_cache_loc=forward_batch.out_cache_loc, out_cache_loc=forward_batch.out_cache_loc,
extend_seq_lens=forward_batch.extend_seq_lens, extend_seq_lens=forward_batch.extend_seq_lens,
seq_lens_cpu=(
seq_lens_cpu.tolist() if seq_lens_cpu is not None else None
),
) )
elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True):
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu