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

Batch EAGLE draft/draft-extend replay memcpys via grouped foreach copy (#28465)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
This commit is contained in:
Khoa Pham
2026-06-17 00:34:21 -07:00
committed by GitHub
co-authored by Cursor Qiaolin Yu
parent 3bc618485a
commit b54f8432ad
2 changed files with 65 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+35 -18
View File
@@ -24,6 +24,7 @@ from sglang.srt.model_executor.runner import (
DecodeCudaGraphRunner,
DeepEPCudaGraphRunnerAdapter,
ShapeKey,
_grouped_foreach_copy_,
get_batch_sizes_to_capture,
model_capture_mode,
)
@@ -459,21 +460,6 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
num_tokens = bs * self.num_tokens_per_bs
# Common inputs
buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
buffers.out_cache_loc[: raw_num_token * self.speculative_num_steps].copy_(
forward_batch.out_cache_loc
)
buffers.positions[:raw_num_token].copy_(forward_batch.positions)
if buffers.rids_int is not None and forward_batch.rids_int is not None:
buffers.rids_int[:raw_bs].copy_(forward_batch.rids_int)
if (
buffers.bootstrap_room_ids_int is not None
and forward_batch.bootstrap_room_ids_int is not None
):
buffers.bootstrap_room_ids_int[:raw_bs].copy_(
forward_batch.bootstrap_room_ids_int
)
maybe_detect_nan(
forward_batch.spec_info.topk_p,
"EagleDraftCudaGraphRunner.replay: topk_p",
@@ -485,14 +471,45 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
"EagleDraftCudaGraphRunner.replay: topk_index vs vocab_size="
f"{self.model_runner.model_config.vocab_size}",
)
buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p)
buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index)
# Common inputs — batch the small per-field device copies into a grouped
# foreach copy (one foreach call per dtype pair) to cut launch overhead.
# hidden_states is handled separately below (see note), and seq_lens_cpu
# is handled further down since it lives on host.
copy_dsts = [
buffers.seq_lens[:raw_bs],
buffers.out_cache_loc[: raw_num_token * self.speculative_num_steps],
buffers.positions[:raw_num_token],
buffers.topk_p[:raw_bs],
buffers.topk_index[:raw_bs],
buffers.req_pool_indices[:raw_bs],
]
copy_srcs = [
forward_batch.seq_lens,
forward_batch.out_cache_loc,
forward_batch.positions,
forward_batch.spec_info.topk_p,
forward_batch.spec_info.topk_index,
forward_batch.req_pool_indices,
]
if buffers.rids_int is not None and forward_batch.rids_int is not None:
copy_dsts.append(buffers.rids_int[:raw_bs])
copy_srcs.append(forward_batch.rids_int)
if (
buffers.bootstrap_room_ids_int is not None
and forward_batch.bootstrap_room_ids_int is not None
):
copy_dsts.append(buffers.bootstrap_room_ids_int[:raw_bs])
copy_srcs.append(forward_batch.bootstrap_room_ids_int)
_grouped_foreach_copy_(copy_dsts, copy_srcs)
# hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync
# DMA engine; foreach would force the ~3x slower compute-kernel copy.
if (
buffers.hidden_states is not None
and forward_batch.spec_info.hidden_states is not None
):
buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
# TODO(ch-wan): support num_token_non_padded
if self.require_gathered_buffer:
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+30 -13
View File
@@ -24,6 +24,7 @@ from sglang.srt.model_executor.runner import (
DecodeCudaGraphRunner,
DeepEPCudaGraphRunnerAdapter,
ShapeKey,
_grouped_foreach_copy_,
get_batch_sizes_to_capture,
model_capture_mode,
)
@@ -455,14 +456,38 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
buffers.num_accept_tokens.fill_(self.num_tokens_per_bs)
buffers.extend_seq_lens.fill_(self.num_tokens_per_bs)
buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids)
buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
# Batch the small per-field device copies into a grouped foreach copy
# (one foreach call per dtype pair) to cut launch overhead. hidden_states
# is handled separately below (see note), and seq_lens_cpu is handled
# further down since it lives on host.
copy_dsts = [
buffers.input_ids[:num_tokens],
buffers.seq_lens[:raw_bs],
buffers.out_cache_loc[:num_tokens],
buffers.positions[:num_tokens],
buffers.req_pool_indices[:raw_bs],
]
copy_srcs = [
forward_batch.input_ids,
forward_batch.seq_lens,
forward_batch.out_cache_loc,
forward_batch.positions,
forward_batch.req_pool_indices,
]
if forward_batch.extend_seq_lens is not None:
buffers.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens)
copy_dsts.append(buffers.extend_seq_lens[:raw_bs])
copy_srcs.append(forward_batch.extend_seq_lens)
else:
buffers.extend_seq_lens[:raw_bs].fill_(self.num_tokens_per_bs)
buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
buffers.positions[:num_tokens].copy_(forward_batch.positions)
if forward_batch.spec_info.num_correct_drafts is not None:
copy_dsts.append(buffers.num_correct_drafts[:raw_bs])
copy_srcs.append(forward_batch.spec_info.num_correct_drafts)
copy_dsts.append(buffers.num_accept_tokens[:raw_bs])
copy_srcs.append(forward_batch.spec_info.num_accept_tokens)
_grouped_foreach_copy_(copy_dsts, copy_srcs)
# hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync
# DMA engine; foreach would force the ~3x slower compute-kernel copy.
if (
buffers.hidden_states is not None
and forward_batch.spec_info.hidden_states is not None
@@ -472,14 +497,6 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
buffers.hidden_states[:num_tokens].copy_(
forward_batch.spec_info.hidden_states
)
if forward_batch.spec_info.num_correct_drafts is not None:
buffers.num_correct_drafts[:raw_bs].copy_(
forward_batch.spec_info.num_correct_drafts
)
buffers.num_accept_tokens[:raw_bs].copy_(
forward_batch.spec_info.num_accept_tokens
)
buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
# TODO(ch-wan): support num_token_non_padded
if self.require_gathered_buffer: