sglang 0.4.1.post2__py3-none-any.whl → 0.4.1.post4__py3-none-any.whl

sglang/srt/managers/tp_worker.py

@@ -30,7 +30,7 @@ from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_a
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
-from sglang.srt.utils import broadcast_pyobj, set_random_seed
+from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
 
 logger = logging.getLogger(__name__)
 
@@ -45,13 +45,18 @@ class TpModelWorker:
         tp_rank: int,
         dp_rank: Optional[int],
         nccl_port: int,
+        is_draft_worker: bool = False,
     ):
         # Parse args
         self.tp_rank = tp_rank
 
         # Init model and tokenizer
         self.model_config = ModelConfig(
-            server_args.model_path,
+            (
+                server_args.model_path
+                if not is_draft_worker
+                else server_args.speculative_draft_model_path
+            ),
             trust_remote_code=server_args.trust_remote_code,
             revision=server_args.revision,
             context_length=server_args.context_length,
@@ -68,6 +73,7 @@ class TpModelWorker:
             tp_size=server_args.tp_size,
             nccl_port=nccl_port,
             server_args=server_args,
+            is_draft_worker=is_draft_worker,
         )
         if server_args.skip_tokenizer_init:
             self.tokenizer = self.processor = None
@@ -150,12 +156,18 @@ class TpModelWorker:
         self,
         model_worker_batch: ModelWorkerBatch,
         launch_done: Optional[threading.Event] = None,
+        skip_sample: bool = False,
     ):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
         logits_output = self.model_runner.forward(forward_batch)
         if launch_done:
             launch_done.set()
-        next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+
+        if skip_sample:
+            next_token_ids = None
+        else:
+            next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+
         return logits_output, next_token_ids
 
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
@@ -191,7 +203,7 @@ class TpModelWorker:
 
     def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
         success, message = self.model_runner.update_weights_from_tensor(
-            recv_req.name, recv_req.tensor
+            MultiprocessingSerializer.deserialize(recv_req.serialized_named_tensors)
         )
         return success, message
 

sglang/srt/managers/tp_worker_overlap_thread.py

@@ -144,10 +144,9 @@ class TpModelWorkerClient:
 
             # Copy results to the CPU
             if model_worker_batch.return_logprob:
-                logits_output.next_token_logprobs = logits_output.next_token_logprobs[
-                    torch.arange(len(next_token_ids), device=self.device),
-                    next_token_ids,
-                ].to("cpu", non_blocking=True)
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.to("cpu", non_blocking=True)
+                )
                 if logits_output.input_token_logprobs is not None:
                     logits_output.input_token_logprobs = (
                         logits_output.input_token_logprobs.to("cpu", non_blocking=True)

sglang/srt/model_executor/cuda_graph_runner.py

@@ -25,14 +25,15 @@ from vllm.distributed import get_tensor_model_parallel_rank
 from vllm.distributed.parallel_state import graph_capture
 from vllm.model_executor.custom_op import CustomOp
 
-from sglang.srt.layers.logits_processor import (
-    LogitsMetadata,
-    LogitsProcessor,
-    LogitsProcessorOutput,
-)
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.moe.fused_moe_native import fused_moe_forward_native
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.utils import maybe_torch_compile, monkey_patch_vllm_all_gather
+from sglang.srt.layers.torchao_utils import save_gemlite_cache
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.utils import monkey_patch_vllm_all_gather
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -105,11 +106,6 @@ def set_torch_compile_config():
         torch._dynamo.config.cache_size_limit = 1024
 
 
-@maybe_torch_compile(dynamic=True)
-def clamp_position(seq_lens):
-    return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
-
-
 class CudaGraphRunner:
     """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
 
@@ -152,6 +148,21 @@ class CudaGraphRunner:
             if bs <= model_runner.req_to_token_pool.size
             and bs <= model_runner.server_args.cuda_graph_max_bs
         ]
+
+        self.capture_forward_mode = ForwardMode.DECODE
+        self.num_tokens_per_bs = 1
+
+        if model_runner.spec_algorithm.is_eagle():
+            if self.model_runner.is_draft_worker:
+                self.num_tokens_per_bs = (
+                    self.model_runner.server_args.speculative_eagle_topk
+                )
+            else:
+                self.capture_forward_mode = ForwardMode.TARGET_VERIFY
+                self.num_tokens_per_bs = (
+                    self.model_runner.server_args.speculative_num_draft_tokens
+                )
+
         self.compile_bs = (
             [
                 bs
@@ -164,8 +175,8 @@ class CudaGraphRunner:
 
         # Attention backend
         self.max_bs = max(self.capture_bs)
-        self.model_runner.attn_backend.init_cuda_graph_state(self.max_bs)
-
+        self.max_num_token = self.max_bs * self.num_tokens_per_bs
+        self.model_runner.attn_backend.init_cuda_graph_state(self.max_num_token)
         self.seq_len_fill_value = (
             self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
         )
@@ -178,14 +189,22 @@ class CudaGraphRunner:
 
         # Common inputs
         with torch.device("cuda"):
-            self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32)
+            self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int32)
             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.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32)
+            self.out_cache_loc = torch.zeros((self.max_num_token,), dtype=torch.int32)
+            self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
             self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int32)
 
+            # Speculative_inference
+            if model_runner.spec_algorithm.is_eagle():
+                self.hidden_states = torch.zeros(
+                    (self.max_num_token, self.model_runner.model_config.hidden_size),
+                    dtype=self.model_runner.dtype,
+                )
+
             if self.is_encoder_decoder:
                 # NOTE: encoder_lens can influence the full_text_row_masked_out_mask tensor when doing mixed batch
                 self.encoder_lens = torch.full(
@@ -257,12 +276,12 @@ class CudaGraphRunner:
     def capture(self):
         with graph_capture() as graph_capture_context:
             self.stream = graph_capture_context.stream
-            capture_bs = (
+            capture_range = (
                 tqdm.tqdm(self.capture_bs)
                 if get_tensor_model_parallel_rank() == 0
                 else self.capture_bs
             )
-            for bs in capture_bs:
+            for bs in capture_range:
                 with patch_model(
                     self.model_runner.model,
                     bs in self.compile_bs,
@@ -276,21 +295,24 @@ class CudaGraphRunner:
                     self.graphs[bs] = graph
                     self.output_buffers[bs] = output_buffers
 
+                # Save gemlite cache after each capture
+                save_gemlite_cache()
+
     def capture_one_batch_size(self, bs: int, forward: Callable):
         graph = torch.cuda.CUDAGraph()
         stream = self.stream
+        num_tokens = bs * self.num_tokens_per_bs
 
         # Common inputs
-        input_ids = self.input_ids[:bs]
+        input_ids = self.input_ids[:num_tokens]
         req_pool_indices = self.req_pool_indices[:bs]
         seq_lens = self.seq_lens[:bs]
-        out_cache_loc = self.out_cache_loc[:bs]
+        out_cache_loc = self.out_cache_loc[:num_tokens]
+        positions = self.positions[:num_tokens]
         if self.is_encoder_decoder:
             encoder_lens = self.encoder_lens[:bs]
         else:
             encoder_lens = None
-
-        seq_lens_sum = seq_lens.sum().item()
         mrope_positions = self.mrope_positions[:, :bs]
 
         if self.enable_dp_attention:
@@ -300,37 +322,43 @@ class CudaGraphRunner:
             global_num_tokens = None
             gathered_buffer = None
 
+        forward_batch = ForwardBatch(
+            forward_mode=self.capture_forward_mode,
+            batch_size=bs,
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            req_to_token_pool=self.model_runner.req_to_token_pool,
+            token_to_kv_pool=self.model_runner.token_to_kv_pool,
+            attn_backend=self.model_runner.attn_backend,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=seq_lens.sum(),
+            encoder_lens=encoder_lens,
+            return_logprob=False,
+            top_logprobs_nums=[0] * bs,
+            positions=positions,
+            global_num_tokens=global_num_tokens,
+            mrope_positions=mrope_positions,
+            gathered_buffer=gathered_buffer,
+            spec_algorithm=self.model_runner.spec_algorithm,
+            spec_info=self.get_spec_info(num_tokens, positions),
+        )
+
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_capture_cuda_graph(
             bs,
+            num_tokens,
             req_pool_indices,
             seq_lens,
             encoder_lens,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
         )
 
         # Run and capture
         def run_once():
-            forward_batch = ForwardBatch(
-                forward_mode=ForwardMode.DECODE,
-                batch_size=bs,
-                input_ids=input_ids,
-                req_pool_indices=req_pool_indices,
-                seq_lens=seq_lens,
-                req_to_token_pool=self.model_runner.req_to_token_pool,
-                token_to_kv_pool=self.model_runner.token_to_kv_pool,
-                attn_backend=self.model_runner.attn_backend,
-                out_cache_loc=out_cache_loc,
-                seq_lens_sum=seq_lens_sum,
-                encoder_lens=encoder_lens,
-                return_logprob=False,
-                top_logprobs_nums=[0] * bs,
-                positions=clamp_position(seq_lens),
-                mrope_positions=mrope_positions,
-                global_num_tokens=global_num_tokens,
-                gathered_buffer=gathered_buffer,
-            )
             logits_output = forward(input_ids, forward_batch.positions, forward_batch)
-            return logits_output.next_token_logits
+            return logits_output.next_token_logits, logits_output.hidden_states
 
         for _ in range(2):
             torch.cuda.synchronize()
@@ -356,6 +384,7 @@ class CudaGraphRunner:
     def replay(self, forward_batch: ForwardBatch):
         assert forward_batch.out_cache_loc is not None
         raw_bs = forward_batch.batch_size
+        raw_num_token = raw_bs * self.num_tokens_per_bs
 
         # Pad
         if self.enable_dp_attention:
@@ -370,15 +399,20 @@ class CudaGraphRunner:
             self.out_cache_loc.zero_()
 
         # Common inputs
-        self.input_ids[:raw_bs].copy_(forward_batch.input_ids)
+        self.input_ids[:raw_num_token].copy_(forward_batch.input_ids)
         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_bs].copy_(forward_batch.out_cache_loc)
+        self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
+        self.positions[:raw_num_token].copy_(forward_batch.positions)
+
         if self.is_encoder_decoder:
             self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)
         if forward_batch.mrope_positions is not None:
             self.mrope_positions[:, :raw_bs].copy_(forward_batch.mrope_positions)
 
+        if hasattr(forward_batch.spec_info, "hidden_states"):
+            self.hidden_states[:raw_num_token] = forward_batch.spec_info.hidden_states
+
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_replay_cuda_graph(
             bs,
@@ -386,40 +420,51 @@ class CudaGraphRunner:
             self.seq_lens,
             forward_batch.seq_lens_sum + (bs - raw_bs),
             self.encoder_lens,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
         )
 
         # Replay
         self.graphs[bs].replay()
-        next_token_logits = self.output_buffers[bs][:raw_bs]
+        next_token_logits, hidden_states = self.output_buffers[bs]
+
+        logits_output = LogitsProcessorOutput(
+            next_token_logits=next_token_logits[:raw_num_token],
+            hidden_states=(
+                hidden_states[:raw_num_token] if hidden_states is not None else None
+            ),
+        )
+        return logits_output
 
-        # Extract logprobs
-        if forward_batch.return_logprob:
-            logits_metadata = LogitsMetadata(
-                forward_mode=ForwardMode.DECODE,
-                top_logprobs_nums=forward_batch.top_logprobs_nums,
+    def get_spec_info(self, num_tokens: int, positions: torch.Tensor):
+        spec_info = None
+        if self.model_runner.spec_algorithm.is_eagle():
+            from sglang.srt.speculative.eagle_utils import (
+                EAGLEDraftInput,
+                EagleVerifyInput,
             )
-            next_token_logprobs = (
-                LogitsProcessor.compute_temp_top_p_normalized_logprobs(
-                    next_token_logits, logits_metadata
+
+            if self.model_runner.is_draft_worker:
+                spec_info = EAGLEDraftInput()
+                spec_info.hidden_states = self.hidden_states[:num_tokens]
+                spec_info.positions = positions
+                spec_info.capture_hidden_mode = CaptureHiddenMode.FULL
+                spec_info.init(self.model_runner.server_args)
+            else:
+                spec_info = EagleVerifyInput(
+                    None,
+                    None,
+                    None,
+                    None,
+                    None,
+                    None,
+                    self.model_runner.server_args.speculative_num_draft_tokens,
                 )
-            )
-            logits_output = LogitsProcessorOutput(
-                next_token_logits=next_token_logits,
-                next_token_logprobs=next_token_logprobs,
-            )
-            return_top_logprob = any(x > 0 for x in forward_batch.top_logprobs_nums)
-            if return_top_logprob:
-                (
-                    logits_output.output_top_logprobs_val,
-                    logits_output.output_top_logprobs_idx,
-                ) = LogitsProcessor.get_top_logprobs(
-                    next_token_logprobs, logits_metadata
-                )[
-                    2:4
-                ]
-        else:
-            logits_output = LogitsProcessorOutput(
-                next_token_logits=next_token_logits,
-            )
+                spec_info.custom_mask = torch.zeros(
+                    (num_tokens * self.model_runner.model_config.context_len),
+                    dtype=torch.bool,
+                    device="cuda",
+                )
+                spec_info.capture_hidden_mode = CaptureHiddenMode.FULL
 
-        return logits_output
+        return spec_info

sglang/srt/model_executor/forward_batch_info.py

@@ -38,6 +38,7 @@ import triton
 import triton.language as tl
 
 from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
+from sglang.srt.utils import maybe_torch_compile
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention import AttentionBackend
@@ -96,7 +97,11 @@ class ForwardMode(IntEnum):
         return self == ForwardMode.DRAFT_EXTEND
 
     def is_cuda_graph(self):
-        return self in (ForwardMode.DECODE, ForwardMode.TARGET_VERIFY)
+        return (
+            self == ForwardMode.DECODE
+            or self == ForwardMode.TARGET_VERIFY
+            or self == ForwardMode.IDLE
+        )
 
     def is_dummy_first(self):
         return self == ForwardMode.DUMMY_FIRST
@@ -161,15 +166,15 @@ class ForwardBatch:
     token_to_kv_pool: BaseTokenToKVPool = None
     attn_backend: AttentionBackend = None
 
-    # Speculative decoding
-    spec_info: SpecInfo = None
-    spec_algorithm: SpeculativeAlgorithm = None
-
     # For DP attention
     global_num_tokens: Optional[List[int]] = None
     gathered_buffer: Optional[torch.Tensor] = None
     can_run_dp_cuda_graph: bool = False
 
+    # Speculative decoding
+    spec_info: SpecInfo = None
+    spec_algorithm: SpeculativeAlgorithm = None
+
     # For Qwen2-VL
     mrope_positions: torch.Tensor = None
 
@@ -258,6 +263,8 @@ class ForwardBatch:
             can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph,
             lora_paths=batch.lora_paths,
             sampling_info=batch.sampling_info,
+            spec_algorithm=batch.spec_algorithm,
+            spec_info=batch.spec_info,
             input_embeds=batch.input_embeds,
         )
 
@@ -270,10 +277,21 @@ class ForwardBatch:
             )
 
         if ret.forward_mode.is_idle():
+            ret.positions = torch.empty((0,), device=device)
             return ret
 
+        # Override the positions with spec_info
+        if (
+            ret.spec_info is not None
+            and getattr(ret.spec_info, "positions", None) is not None
+        ):
+            ret.positions = ret.spec_info.positions
+
         # Init position information
-        if not ret.forward_mode.is_decode():
+        if ret.forward_mode.is_decode():
+            if ret.positions is None:
+                ret.positions = clamp_position(batch.seq_lens)
+        else:
             ret.extend_seq_lens = torch.tensor(
                 batch.extend_seq_lens, dtype=torch.int32
             ).to(device, non_blocking=True)
@@ -282,13 +300,15 @@ class ForwardBatch:
             ).to(device, non_blocking=True)
             if model_runner.server_args.attention_backend != "torch_native":
                 ret.extend_num_tokens = batch.extend_num_tokens
-                ret.positions, ret.extend_start_loc = compute_position_triton(
+                positions, ret.extend_start_loc = compute_position_triton(
                     ret.extend_prefix_lens, ret.extend_seq_lens, ret.extend_num_tokens
                 )
             else:
-                ret.positions, ret.extend_start_loc = compute_position_torch(
+                positions, ret.extend_start_loc = compute_position_torch(
                     ret.extend_prefix_lens, ret.extend_seq_lens
                 )
+            if ret.positions is None:
+                ret.positions = positions
             ret.extend_prefix_lens_cpu = batch.extend_prefix_lens
             ret.extend_seq_lens_cpu = batch.extend_seq_lens
             ret.extend_logprob_start_lens_cpu = batch.extend_logprob_start_lens
@@ -377,6 +397,11 @@ def compute_position_torch(
     return positions.to(torch.int64), extend_start_loc
 
 
+@maybe_torch_compile(dynamic=True)
+def clamp_position(seq_lens):
+    return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
+
+
 class CaptureHiddenMode(IntEnum):
     NULL = auto()
     FULL = auto()