diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 3dc6c11dc..1c5a153a1 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -802,7 +802,7 @@ class CPUGraphRunner: captured_forward_batch.encoder_lens[:raw_bs].copy_( forward_batch.encoder_lens ) - if enable_num_token_non_padded(self.model_runner.server_args): + if enable_num_token_non_padded(): captured_forward_batch.num_token_non_padded.copy_( forward_batch.num_token_non_padded ) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index b04ee47cd..e5e6b02bd 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -987,7 +987,7 @@ class CudaGraphRunner: # populate_from_forward_batch). buffers.num_token_non_padded[...] = num_tokens if ( - enable_num_token_non_padded(self.model_runner.server_args) + enable_num_token_non_padded() and self.require_gathered_buffer and not self.nsa_enable_prefill_cp ): @@ -1255,9 +1255,7 @@ class CudaGraphRunner: require_gathered_buffer=self.require_gathered_buffer, num_tokens_per_bs=self.num_tokens_per_bs, nsa_enable_prefill_cp=self.nsa_enable_prefill_cp, - enable_num_token_non_padded_flag=enable_num_token_non_padded( - self.model_runner.server_args - ), + enable_num_token_non_padded_flag=enable_num_token_non_padded(), pp_proxy_tensors=pp_proxy_tensors, ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 3192138f7..7ea1336c3 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -499,7 +499,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ) num_tokens = len(batch.input_ids) if batch.input_ids is not None else 0 - if enable_num_token_non_padded(model_runner.server_args): + if enable_num_token_non_padded(): ret.num_token_non_padded = torch.tensor(num_tokens, dtype=torch.int32).to( device, non_blocking=True ) @@ -1086,7 +1086,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): return self.tbo_split_seq_index is not None -def enable_num_token_non_padded(server_args): +def enable_num_token_non_padded(): return get_moe_expert_parallel_world_size() > 1 diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index a5ae5b5b3..f477f4ef9 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -365,7 +365,6 @@ class EAGLEDraftExtendCudaGraphRunner: num_accepted_drafts=num_accepted_drafts, num_accepted_tokens=num_accepted_tokens, ) - spec_info.positions = None self.deepep_adapter.capture(is_extend_in_batch=True) diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 3c480e9a1..751dee728 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -485,8 +485,8 @@ class EAGLEWorker(TpModelWorker): set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) - logits_output, verify_output, model_worker_batch, can_run_cuda_graph = ( - self.verify(batch, spec_info) + logits_output, verify_output, can_run_cuda_graph = self.verify( + batch, spec_info ) if get_global_tracing_enabled(): @@ -527,9 +527,7 @@ class EAGLEWorker(TpModelWorker): can_run_cuda_graph=can_run_cuda_graph, ) - def check_forward_draft_extend_after_decode( - self, batch: ScheduleBatch, verify_output: EagleVerifyOutput - ): + def check_forward_draft_extend_after_decode(self, verify_output: EagleVerifyOutput): local_need_forward = verify_output.unfinished_accept_tokens.shape[0] > 0 if not self.server_args.enable_dp_attention: return local_need_forward @@ -981,7 +979,7 @@ class EAGLEWorker(TpModelWorker): ) batch.spec_info = res.next_draft_input - return logits_output, res, model_worker_batch, can_run_cuda_graph + return logits_output, res, can_run_cuda_graph def _mamba_verify_update( self, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index 9039577cc..09ea8f0b9 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -447,9 +447,7 @@ class FrozenKVMTPWorker(TpModelWorker): set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) - logits_output, verify_output, _, can_run_cuda_graph = self.verify( - batch, spec_info - ) + logits_output, verify_output, can_run_cuda_graph = self.verify(batch, spec_info) if get_global_tracing_enabled(): for idx, req in enumerate(batch.reqs): @@ -771,4 +769,4 @@ class FrozenKVMTPWorker(TpModelWorker): batch.spec_info = res.next_draft_input del seq_lens_pre_verify - return logits_output, res, model_worker_batch, can_run_cuda_graph + return logits_output, res, can_run_cuda_graph diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index b03492905..8dcb6685c 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -272,8 +272,8 @@ class MultiLayerEagleWorker(TpModelWorker): self.mtp_model_runner(0).tp_group ), speculative_moe_backend_context(): spec_info = self.draft(batch) - logits_output, verify_output, model_worker_batch, can_run_cuda_graph = ( - self.verify(batch, spec_info) + logits_output, verify_output, can_run_cuda_graph = self.verify( + batch, spec_info ) with self.draft_tp_context( @@ -295,9 +295,7 @@ class MultiLayerEagleWorker(TpModelWorker): can_run_cuda_graph=can_run_cuda_graph, ) - def check_forward_draft_extend_after_decode( - self, batch: ScheduleBatch, verify_output: EagleVerifyOutput - ): + def check_forward_draft_extend_after_decode(self, verify_output: EagleVerifyOutput): local_need_forward = verify_output.unfinished_accept_tokens.shape[0] > 0 if not self.server_args.enable_dp_attention: return local_need_forward @@ -593,7 +591,7 @@ class MultiLayerEagleWorker(TpModelWorker): ) batch.spec_info = res.next_draft_input - return logits_output, res, model_worker_batch, can_run_cuda_graph + return logits_output, res, can_run_cuda_graph def forward_draft_extend( self,