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

align code style eagle draft&draft_extend cuda graph runner (#13533)

Browse Source
This commit is contained in:
yinghui
2025-11-22 18:55:16 +01:00
committed by GitHub
parent dd303614e0
commit 5a4394a342
2 changed files with 19 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
View File

@@ -33,10 +33,6 @@ from sglang.srt.utils import (
if TYPE_CHECKING:
from sglang.srt.speculative.eagle_worker import EAGLEWorker
import logging
logger = logging.getLogger(__name__)
class EAGLEDraftCudaGraphRunner:
def __init__(self, eagle_worker: EAGLEWorker):
@@ -51,29 +47,28 @@ class EAGLEDraftCudaGraphRunner:
self.output_buffers = {}
self.enable_torch_compile = model_runner.server_args.enable_torch_compile
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
self.dp_size = self.model_runner.dp_size
self.tp_size = self.model_runner.tp_size
self.topk = model_runner.server_args.speculative_eagle_topk
self.dp_size = self.model_runner.dp_size
self.speculative_num_steps = model_runner.server_args.speculative_num_steps
self.topk = model_runner.server_args.speculative_eagle_topk
self.enable_profile_cuda_graph = (
model_runner.server_args.enable_profile_cuda_graph
)
self.enable_pdmux = False
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
server_args = model_runner.server_args
# Batch sizes to capture
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
self.num_tokens_per_bs = server_args.speculative_eagle_topk
# Attention backend
self.num_tokens_per_bs = self.topk
self.max_bs = max(self.capture_bs)
self.max_num_token = self.max_bs * self.num_tokens_per_bs
self.model_runner.draft_attn_backend.init_cuda_graph_state(
self.max_bs, self.max_num_token
)
@@ -92,10 +87,6 @@ class EAGLEDraftCudaGraphRunner:
with torch.device(model_runner.device):
self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
self.seq_lens = torch.full(
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
)
self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
self.out_cache_loc = torch.zeros(
(self.max_num_token * self.speculative_num_steps,), dtype=torch.int64
)
@@ -103,6 +94,10 @@ class EAGLEDraftCudaGraphRunner:
self.mrope_positions = torch.zeros(
(3, self.max_num_token), dtype=torch.int64
)
self.seq_lens = torch.full(
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
)
self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32)
self.topk_p = torch.zeros((self.max_bs, self.topk), dtype=torch.float32)
self.topk_index = torch.zeros((self.max_bs, self.topk), dtype=torch.int64)
self.hidden_states = torch.zeros(
@@ -194,9 +189,9 @@ class EAGLEDraftCudaGraphRunner:
out_cache_loc = self.out_cache_loc[: num_tokens * self.speculative_num_steps]
positions = self.positions[:num_tokens]
mrope_positions = self.mrope_positions[:, :num_tokens]
hidden_states = self.hidden_states[:num_seqs]
topk_p = self.topk_p[:num_seqs]
topk_index = self.topk_index[:num_seqs]
hidden_states = self.hidden_states[:num_seqs]
if self.require_mlp_tp_gather:
self.global_num_tokens_gpu.copy_(
@@ -264,6 +259,7 @@ class EAGLEDraftCudaGraphRunner:
positions=positions,
mrope_positions=mrope_positions,
global_num_tokens_gpu=global_num_tokens,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=global_dp_buffer_len,
spec_algorithm=self.model_runner.spec_algorithm,
@@ -271,7 +267,6 @@ class EAGLEDraftCudaGraphRunner:
capture_hidden_mode=(
spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
),
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob,
)
# Attention backend
@@ -303,6 +298,7 @@ class EAGLEDraftCudaGraphRunner:
self.deepep_adapter.capture(is_extend_in_batch=False)
self._capture_init(run_once)
out = self._capture_graph(
graph, get_global_graph_memory_pool(), stream, run_once
)
@@ -333,6 +329,7 @@ class EAGLEDraftCudaGraphRunner:
index = bisect.bisect_left(self.capture_bs, max_batch_size)
else:
index = bisect.bisect_left(self.capture_bs, raw_bs)
bs = self.capture_bs[index]
if bs != raw_bs:
self.seq_lens.fill_(self.seq_len_fill_value)
@@ -342,7 +339,6 @@ class EAGLEDraftCudaGraphRunner:
num_tokens = bs * self.num_tokens_per_bs
# Common inputs
self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
self.out_cache_loc[: raw_num_token * self.speculative_num_steps].copy_(
forward_batch.out_cache_loc
@@ -351,6 +347,7 @@ class EAGLEDraftCudaGraphRunner:
self.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p)
self.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index)
self.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
# TODO(ch-wan): support num_token_non_padded
if self.require_gathered_buffer:

11
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
View File

@@ -57,16 +57,17 @@ class EAGLEDraftExtendCudaGraphRunner:
self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
self.tp_size = self.model_runner.tp_size
self.dp_size = model_runner.server_args.dp_size
self.dp_size = self.model_runner.dp_size
self.speculative_num_steps = model_runner.server_args.speculative_num_steps
self.topk = model_runner.server_args.speculative_eagle_topk
self.enable_profile_cuda_graph = (
model_runner.server_args.enable_profile_cuda_graph
)
self.enable_pdmux = False
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
self.padded_static_len = -1
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
# Attention backend
self.num_tokens_per_bs = self.speculative_num_steps + 1
@@ -232,11 +233,11 @@ class EAGLEDraftExtendCudaGraphRunner:
seq_lens_cpu = self.seq_lens_cpu[:bs]
extend_seq_lens = self.extend_seq_lens[:bs]
extend_seq_lens_cpu = self.extend_seq_lens_cpu[:bs]
accept_length = self.accept_length[:bs]
out_cache_loc = self.out_cache_loc[:num_tokens]
positions = self.positions[:num_tokens]
mrope_positions = self.mrope_positions[:, :num_tokens]
hidden_states = self.hidden_states[:num_tokens]
accept_length = self.accept_length[:bs]
next_token_logits_buffer = self.next_token_logits_buffer[
: bs if self.forward_mode == ForwardMode.DRAFT_EXTEND else num_tokens
]
@@ -293,6 +294,8 @@ class EAGLEDraftExtendCudaGraphRunner:
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu,
next_token_logits_buffer=next_token_logits_buffer,
extend_seq_lens=extend_seq_lens,
extend_seq_lens_cpu=extend_seq_lens_cpu,
req_to_token_pool=self.model_runner.req_to_token_pool,
token_to_kv_pool=self.model_runner.token_to_kv_pool,
out_cache_loc=out_cache_loc,
@@ -308,8 +311,6 @@ class EAGLEDraftExtendCudaGraphRunner:
spec_info=spec_info,
capture_hidden_mode=CaptureHiddenMode.LAST,
attn_backend=self.eagle_worker.draft_extend_attn_backend,
extend_seq_lens=extend_seq_lens,
extend_seq_lens_cpu=extend_seq_lens_cpu,
padded_static_len=self.padded_static_len,
)