# Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import contextlib import logging from typing import TYPE_CHECKING, List, Optional, Tuple import torch from sglang.srt.environ import envs from sglang.srt.layers.moe.utils import speculative_moe_backend_context from sglang.srt.managers.schedule_batch import ModelWorkerBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.eagle_info_v2 import fill_new_verified_id from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, ) from sglang.srt.speculative.multi_layer_eagle_utils import ( assign_hidden_states_pool_triton, rotate_input_ids_triton, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( detect_nan, draft_tp_context, select_top_k_tokens, ) from sglang.srt.utils.common import empty_context, fast_topk if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput logger = logging.getLogger(__name__) def _get_plan_stream( device: str, ) -> Tuple[any, contextlib.AbstractContextManager]: if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get(): plan_stream = torch.get_device_module(device).Stream() plan_stream_ctx = torch.get_device_module(device).stream(plan_stream) return plan_stream, plan_stream_ctx else: return None, contextlib.nullcontext() class MultiLayerEagleDraftWorker(BaseDraftWorker): def __init__( self, server_args: ServerArgs, gpu_id: int, tp_rank: int, dp_rank: int, moe_ep_rank: int, attn_cp_rank: int, moe_dp_rank: int, nccl_port: int, target_worker: TpModelWorker, ): # copy args self.server_args = server_args self.gpu_id = gpu_id self.tp_rank = tp_rank self.dp_rank = dp_rank self.moe_ep_rank = moe_ep_rank self.nccl_port = nccl_port self.target_worker = target_worker self.draft_extend_attn_backend_list = [] self.model_config = target_worker.model_config # Args for easy access self.device = server_args.device self.topk = server_args.speculative_eagle_topk self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) # Set constant EagleDraftInput.ALLOC_LEN_PER_DECODE = max( self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens ) # Do not capture cuda graph in `TpModelWorker` init, # will capture later with init_cuda_graphs() backup_disable_cuda_graph = server_args.disable_cuda_graph server_args.disable_cuda_graph = True # Share the allocator with a target worker. # Draft and target worker own their own KV cache pools. self.req_to_token_pool, self.token_to_kv_pool_allocator = ( target_worker.get_memory_pool() ) with empty_context(), speculative_moe_backend_context(): # Init draft worker self.draft_worker = TpModelWorker( server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, pp_rank=0, # FIXME dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, moe_dp_rank=moe_dp_rank, nccl_port=nccl_port, is_draft_worker=True, req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, is_multi_layer_eagle=True, ) # Alias for better readability self.draft_runner_list: List[ModelRunner] = self.draft_worker.model_runner_list self.init_lm_head() # Used for KV Cache reversion self.req_to_hidden_states_pool = torch.empty( ( self.req_to_token_pool.size, self.speculative_num_steps - 1, self.model_config.hidden_size, ), dtype=self.model_config.dtype, device=self.device, ) # Init attention backend and cuda graphs for i in range(self.speculative_num_steps): self.draft_runner_list[i].server_args.disable_cuda_graph = ( backup_disable_cuda_graph ) self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context ) with self.draft_tp_context( self.draft_runner_list[0].tp_group ), speculative_moe_backend_context(): self.init_attention_backend() self.init_cuda_graphs() self.tree_mask_mode = TreeMaskMode.FULL_MASK self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) def mtp_model_runner(self, step: int): return self.draft_runner_list[step] def init_lm_head(self): embed, head = self.target_worker.model_runner.model.get_embed_and_head() # Share the embedding and lm_head for i in range(self.speculative_num_steps): self.draft_runner_list[i].model.set_embed_and_head(embed, head) def init_attention_backend(self): # Create attn backends self.draft_extend_attn_backend_list = [] for step in range(self.speculative_num_steps): from sglang.srt.layers.attention.flashattention_backend import ( FlashAttentionBackend, ) self.draft_extend_attn_backend_list.append( FlashAttentionBackend( model_runner=self.draft_runner_list[step], skip_prefill=False, speculative_step_id=step, ) ) self.draft_runner_list[step].attn_backend = ( self.draft_extend_attn_backend_list[-1] ) def init_cuda_graphs(self): """Capture cuda graphs.""" self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None if self.server_args.disable_cuda_graph: return self.cuda_graph_runner_for_draft_extend = ( MultiLayerEagleMultiStepDraftExtendCudaGraphRunner(self) ) def reset_cuda_graph_buffers(self, forward_batch, batch_result): if self.cuda_graph_runner_for_draft_extend: self.cuda_graph_runner_for_draft_extend.reset_buffers( forward_batch, batch_result ) def draft(self, model_worker_batch: ModelWorkerBatch): draft_input: EagleDraftInput = model_worker_batch.spec_info forward_batch, can_cuda_graph = draft_input.prepare_for_v2_draft( self.req_to_token_pool, model_worker_batch, self.cuda_graph_runner, self.draft_runner_list[0], self.topk, self.speculative_num_steps, ) # Run draft parent_list, top_scores_index, draft_tokens = self.draft_forward(forward_batch) if model_worker_batch.forward_mode.is_idle(): return EagleVerifyInput.create_idle_input( self.topk, self.speculative_num_steps, self.speculative_num_draft_tokens, ) # Build tree mask # Directly write to cuda graph buffers for verify attn tree_mask_buf, position_buf = ( self.target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft() ) ( tree_mask, position, retrive_index, retrive_next_token, retrive_next_sibling, draft_tokens, ) = build_tree_kernel_efficient( draft_input.verified_id, parent_list, top_scores_index, draft_tokens, model_worker_batch.seq_lens, model_worker_batch.seq_lens_sum, self.topk, self.speculative_num_steps, self.speculative_num_draft_tokens, self.tree_mask_mode, tree_mask_buf, position_buf, ) return EagleVerifyInput( draft_token=draft_tokens, custom_mask=tree_mask, positions=position, retrive_index=retrive_index, retrive_next_token=retrive_next_token, retrive_next_sibling=retrive_next_sibling, retrive_cum_len=None, spec_steps=self.speculative_num_steps, topk=self.topk, draft_token_num=self.speculative_num_draft_tokens, capture_hidden_mode=None, seq_lens_sum=None, seq_lens_cpu=None, ) def draft_forward(self, forward_batch: ForwardBatch): # Parse args spec_info: EagleDraftInput = forward_batch.spec_info topk_p, topk_index, hidden_states = ( spec_info.topk_p, spec_info.topk_index, spec_info.hidden_states, ) # Return values score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] parents_list: List[torch.Tensor] = [] # Forward multiple steps scores = None _, hidden_states, scores, tree_info = select_top_k_tokens( 0, topk_p, topk_index, hidden_states, scores, self.topk ) if self.speculative_num_steps == 1: score_list.append(tree_info[0]) token_list.append(tree_info[1]) parents_list.append(tree_info[2]) else: for i in range(self.speculative_num_steps): score_list.append(tree_info[0][:, :, i].unsqueeze(-1)) token_index = tree_info[1][:, i].unsqueeze(-1) token_list.append(token_index) if i == 0: parents_list.append(tree_info[2]) else: parents_list.append( torch.full( (tree_info[2].size(0), 1), i, dtype=torch.long, device="cuda", ) ) # Organize the results score_list = torch.cat(score_list, dim=1).flatten( 1 ) # b, n, topk; n= 1 + (num_steps-1) * self.topk ss_token_list = torch.cat( token_list, dim=1 ) # b, (self.topk + (num_steps-1) * self.topk) top_scores = torch.topk( score_list, self.speculative_num_draft_tokens - 1, dim=-1 ) top_scores_index = top_scores.indices top_scores_index = torch.sort(top_scores_index).values draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) if len(parents_list) > 1: parent_list = torch.cat(parents_list[:-1], dim=1) else: batch_size = parents_list[0].shape[0] parent_list = torch.empty(batch_size, 0, device=parents_list[0].device) return parent_list, top_scores_index, draft_tokens def draft_extend(self): pass def _draft_extend_for_prefill( self, batch: ModelWorkerBatch, target_hidden_states: torch.Tensor, next_token_ids: torch.Tensor, ): """ Run draft model extend to correctly fill the KV cache. Args: batch: The batch to run. target_hidden_states: Hidden states from the target model forward next_token_ids: Next token ids generated from the target forward. """ # Construct spec_info next_draft_input = EagleDraftInput( hidden_states=target_hidden_states, verified_id=next_token_ids, new_seq_lens=batch.seq_lens, # draft mode is same with decode mode, only 1 token per req num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, ) batch.spec_info = next_draft_input # Run forward forward_batch = ForwardBatch.init_new(batch, self.draft_runner_list[0]) forward_batch.return_hidden_states_before_norm = True # Construct input_ids if not batch.forward_mode.is_idle(): rotate_input_ids_triton( forward_batch.input_ids, forward_batch.extend_start_loc, forward_batch.extend_seq_lens, next_token_ids, ) topk_p_list = [] topk_index_list = [] for step in range(self.speculative_num_steps): output: ModelRunnerOutput = self.draft_runner_list[step].forward( forward_batch ) probs = torch.softmax(output.logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) topk_index_list.append(topk_index) if forward_batch.extend_seq_lens is not None: rotate_input_ids_triton( forward_batch.input_ids, forward_batch.extend_start_loc, forward_batch.extend_seq_lens, topk_index, ) next_draft_input.topk_p = torch.cat(topk_p_list, dim=1) next_draft_input.topk_index = torch.cat(topk_index_list, dim=1) # Update req_to_hidden_states_pool for KV Cache reversion if forward_batch.extend_seq_lens is not None: assign_hidden_states_pool_triton( target_hidden_states, forward_batch.req_pool_indices, self.req_to_hidden_states_pool, self.speculative_num_steps - 1, forward_batch.batch_size, forward_batch.extend_seq_lens, forward_batch.extend_start_loc, ) return next_draft_input def _draft_extend_for_decode( self, batch: ModelWorkerBatch, batch_result: GenerationBatchResult ): # Batch 2: Draft extend draft_input = EagleDraftInput( hidden_states=batch_result.logits_output.hidden_states, num_tokens_per_req=self.speculative_num_steps + 1, num_tokens_for_logprob_per_req=1, ) # Prepare for draft extend in a separate stream # Notice that here we use batch_result.next_token_ids as the input ids with self.plan_stream_ctx: forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache( batch, batch_result.next_token_ids, self.speculative_num_draft_tokens, self.draft_runner_list[0], self.cuda_graph_runner_for_draft_extend, ) forward_batch.return_hidden_states_before_norm = True if self.plan_stream: torch.get_device_module(self.device).current_stream().wait_stream( self.plan_stream ) # Run draft extend batch in the main compute stream can_cuda_graph = ( self.cuda_graph_runner_for_draft_extend and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) ) ret_topk_p_list = [] ret_topk_index_list = [] next_token_ids_backup = batch_result.next_token_ids.clone() if can_cuda_graph: self.reset_cuda_graph_buffers(forward_batch, batch_result) else: logger.warning_once( f"can't use cuda graph for draft extend! may have correctness issue!" ) select_index = ( torch.arange(len(batch.seq_lens), device=self.device) * self.speculative_num_draft_tokens + batch_result.accept_lens - 1 ) for step in range(self.speculative_num_steps): # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}") if can_cuda_graph: draft_logits_output = ( self.cuda_graph_runner_for_draft_extend.get_runner(step).replay( forward_batch, init_state=(step == 0) ) ) ret_topk_p, ret_topk_index = ( draft_logits_output.topk_p, draft_logits_output.topk_index, ) else: draft_logits_output = self.draft_runner_list[step].forward( forward_batch, skip_attn_backend_init=True ) probs = torch.softmax( draft_logits_output.logits_output.next_token_logits[select_index], dim=-1, ) ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) if forward_batch.extend_seq_lens is not None: rotate_input_ids_triton( forward_batch.input_ids, forward_batch.extend_start_loc, forward_batch.extend_seq_lens, ret_topk_index, select_index, ) ret_topk_p_list.append(ret_topk_p) ret_topk_index_list.append(ret_topk_index) # Update req_to_hidden_states_pool for KV Cache reversion if ( self.cuda_graph_runner_for_draft_extend is not None and forward_batch.extend_seq_lens is not None ): last_cuda_graph_runner = ( self.cuda_graph_runner_for_draft_extend.get_last_runner() ) assign_hidden_states_pool_triton( last_cuda_graph_runner.hidden_states, last_cuda_graph_runner.req_pool_indices, self.req_to_hidden_states_pool, self.speculative_num_steps - 1, forward_batch.batch_size, last_cuda_graph_runner.extend_seq_lens, last_cuda_graph_runner.extend_start_loc, ) # Reorganize the spec info for the next batch # draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[ # select_index # ] # draft_logits_output.hidden_states = draft_logits_output.hidden_states[ # select_index # ] batch_result.next_token_ids = next_token_ids_backup # Construct the return values next_draft_input = batch_result.next_draft_input ( next_draft_input.topk_p, next_draft_input.topk_index, next_draft_input.hidden_states, ) = ( torch.cat(ret_topk_p_list, dim=1).clone(), torch.cat(ret_topk_index_list, dim=1).clone(), None, ) class MultiLayerEagleWorkerV2(BaseSpecWorker): def __init__( self, server_args: ServerArgs, gpu_id: int, tp_rank: int, dp_rank: Optional[int], moe_ep_rank: int, attn_cp_rank: int, moe_dp_rank: int, nccl_port: int, target_worker: TpModelWorker, ): # Parse arguments self.server_args = server_args self.topk = server_args.speculative_eagle_topk self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens self.enable_nan_detection = server_args.enable_nan_detection self.gpu_id = gpu_id self.device = server_args.device self._target_worker = target_worker self.page_size = server_args.page_size self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) self.req_to_token_pool, self.token_to_kv_pool_allocator = ( target_worker.get_memory_pool() ) # Override the context length of the draft model to be the same as the target model. server_args.context_length = target_worker.model_runner.model_config.context_len self._draft_worker = MultiLayerEagleDraftWorker( server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, attn_cp_rank, moe_dp_rank, nccl_port, target_worker, ) # Some dummy tensors self.num_new_pages_per_topk = torch.empty( (), dtype=torch.int64, device=self.device ) self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device) self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) @property def target_worker(self): return self._target_worker @property def draft_worker(self): return self._draft_worker def clear_cache_pool(self): # allocator and kv cache pool are shared with target worker, which are cleared in scheduler pass def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch): if ( model_worker_batch.forward_mode.is_extend() or model_worker_batch.is_extend_in_batch ): # Target prefill model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL batch_output = self.target_worker.forward_batch_generation( model_worker_batch ) # Draft prefill model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill( model_worker_batch, batch_output.logits_output.hidden_states, batch_output.next_token_ids, ) return batch_output else: if model_worker_batch.spec_info is None: model_worker_batch.spec_info = EagleDraftInput.create_idle_input( device=self.device, hidden_size=self.target_worker.model_config.hidden_size, dtype=self.target_worker.model_config.dtype, topk=self.topk * self.speculative_num_steps, capture_hidden_mode=CaptureHiddenMode.LAST, ) draft_input: EagleDraftInput = model_worker_batch.spec_info verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch) assert verify_input.is_verify_input() model_worker_batch.spec_info = verify_input batch_output = self.verify(model_worker_batch) self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output) return batch_output def verify( self, batch: ModelWorkerBatch, ): # Since batch.seq_lens is allocated in another stream, we need # record_stream() to prevent pytorch gc and reuse the gpu memory # while forward_stream is still running. batch.seq_lens.record_stream( torch.get_device_module(self.device).current_stream() ) # Parse args verify_input: EagleVerifyInput = batch.spec_info bs = len(batch.seq_lens) # Batch 1: Target verify # Prepare for target verify in a separate stream with self.plan_stream_ctx: verify_forward_batch, can_run_cuda_graph = ( verify_input.prepare_for_v2_verify( self.req_to_token_pool, batch, self.target_worker, ) ) # Correct some buffers due to the overlap plan if self.plan_stream: torch.get_device_module(self.device).current_stream().wait_stream( self.plan_stream ) # Some values such as custom_mask and position depend on the output of draft, # so the previous plan step used the wrong values. Here, we need to run the related # computation again to update them to the correct values. self.target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft( verify_input, ( self.target_worker.model_runner.graph_runner.bs if can_run_cuda_graph else None ), ) # Run target verify batch in the main compute stream forward_batch_output = self.target_worker.forward_batch_generation( model_worker_batch=None, forward_batch=verify_forward_batch, is_verify=True, skip_attn_backend_init=True, ) logits_output = forward_batch_output.logits_output # Sample if self.enable_nan_detection: detect_nan(logits_output) ( predict, accept_length, accept_index, ) = verify_input.sample(batch, logits_output) new_seq_lens = batch.seq_lens + accept_length verify_done = torch.get_device_module(self.device).Event() verify_done.record() if not batch.forward_mode.is_idle(): all_verified_id = predict[accept_index] verified_id = torch.empty_like(accept_length, dtype=torch.int32) fill_new_verified_id[(bs,)]( all_verified_id, accept_length, verified_id, self.speculative_num_draft_tokens, ) else: verified_id = torch.empty((0,), device=self.device, dtype=torch.int32) # Construct the next draft input next_draft_input = EagleDraftInput( verified_id=verified_id, new_seq_lens=new_seq_lens, verify_done=verify_done, ) return GenerationBatchResult( logits_output=logits_output, next_token_ids=predict, can_run_cuda_graph=can_run_cuda_graph, next_draft_input=next_draft_input, accept_lens=accept_length, )