sglang 0.4.3.post3__py3-none-any.whl → 0.4.3.post4__py3-none-any.whl

sglang/srt/model_executor/model_runner.py

@@ -122,66 +122,17 @@ class ModelRunner:
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
 
         # Model-specific adjustment
-        if (
-            self.model_config.attention_arch == AttentionArch.MLA
-            and not self.server_args.disable_mla
-        ):
-            # TODO: add MLA optimization on CPU
-            if self.server_args.device != "cpu":
-                if server_args.enable_flashinfer_mla:
-                    logger.info(
-                        "MLA optimization is turned on. Use flashinfer mla backend."
-                    )
-                    self.server_args.attention_backend = "flashinfer_mla"
-                else:
-                    logger.info("MLA optimization is turned on. Use triton backend.")
-                    self.server_args.attention_backend = "triton"
+        self.model_specific_adjustment()
 
-        if self.server_args.enable_double_sparsity:
-            logger.info(
-                "Double sparsity optimization is turned on. Use triton backend without CUDA graph."
-            )
-            self.server_args.attention_backend = "triton"
-            self.server_args.disable_cuda_graph = True
-            if self.server_args.ds_heavy_channel_type is None:
-                raise ValueError(
-                    "Please specify the heavy channel type for double sparsity optimization."
-                )
-            self.init_double_sparsity_channel_config(
-                self.server_args.ds_heavy_channel_type
-            )
-
-        if self.is_multimodal:
-            self.mem_fraction_static *= 0.95
-            logger.info(
-                f"Automatically reduce --mem-fraction-static to {self.mem_fraction_static:.3f} "
-                f"because this is a multimodal model."
-            )
-
-            if self.model_config.hf_config.architectures == [
-                "MllamaForConditionalGeneration"
-            ]:
-                logger.info("Automatically turn off --chunked-prefill-size for mllama.")
-                server_args.chunked_prefill_size = -1
-
-            if self.model_config.hf_config.architectures == [
-                "Qwen2VLForConditionalGeneration"
-            ]:
-                # TODO: qwen2-vl does not support radix cache now, set disable_radix_cache=True automatically
-                logger.info(
-                    "Automatically turn off --chunked-prefill-size and disable radix cache for qwen2-vl."
-                )
-                server_args.chunked_prefill_size = -1
-                server_args.disable_radix_cache = True
-
-        # Global vars
         if server_args.show_time_cost:
             enable_show_time_cost()
+
         if server_args.disable_outlines_disk_cache:
             from outlines.caching import disable_cache
 
             disable_cache()
 
+        # Global vars
         global_server_args_dict.update(
             {
                 "attention_backend": server_args.attention_backend,
@@ -203,6 +154,7 @@ class ModelRunner:
             }
         )
 
+        # CPU offload
         set_cpu_offload_max_bytes(int(server_args.cpu_offload_gb * 1024**3))
 
         # Get memory before model loading
@@ -216,18 +168,6 @@ class ModelRunner:
         self.sampler = Sampler()
         self.load_model()
 
-        # Handle the case where some of models don't finish loading.
-        try:
-            dist.monitored_barrier(
-                group=get_tp_group().cpu_group,
-                timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S),
-                wait_all_ranks=True,
-            )
-        except RuntimeError:
-            raise ValueError(
-                f"TP rank {self.tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node."
-            ) from None
-
         # Apply torchao quantization
         torchao_applied = getattr(self.model, "torchao_applied", False)
         # In layered loading, torchao may have been applied
@@ -244,9 +184,11 @@ class ModelRunner:
         else:
             self.torch_tp_applied = False
 
-        # Init memory pool and attention backends
+        # Init lora
         if server_args.lora_paths is not None:
             self.init_lora_manager()
+
+        # Init memory pool and attention backends
         self.init_memory_pool(
             min_per_gpu_memory,
             server_args.max_running_requests,
@@ -260,10 +202,63 @@ class ModelRunner:
             self.cuda_graph_runner = None
             self.init_attention_backend()
 
+    def model_specific_adjustment(self):
+        server_args = self.server_args
+
+        if (
+            self.model_config.attention_arch == AttentionArch.MLA
+            and not server_args.disable_mla
+        ):
+            # TODO: add MLA optimization on CPU
+            if server_args.device != "cpu":
+                if server_args.enable_flashinfer_mla:
+                    logger.info(
+                        "MLA optimization is turned on. Use flashinfer mla backend."
+                    )
+                    server_args.attention_backend = "flashinfer_mla"
+                else:
+                    logger.info("MLA optimization is turned on. Use triton backend.")
+                    server_args.attention_backend = "triton"
+
+        if server_args.enable_double_sparsity:
+            logger.info(
+                "Double sparsity optimization is turned on. Use triton backend without CUDA graph."
+            )
+            server_args.attention_backend = "triton"
+            server_args.disable_cuda_graph = True
+            if server_args.ds_heavy_channel_type is None:
+                raise ValueError(
+                    "Please specify the heavy channel type for double sparsity optimization."
+                )
+            self.init_double_sparsity_channel_config(server_args.ds_heavy_channel_type)
+
+        if self.is_multimodal:
+            self.mem_fraction_static *= 0.95
+            logger.info(
+                f"Automatically reduce --mem-fraction-static to {self.mem_fraction_static:.3f} "
+                f"because this is a multimodal model."
+            )
+
+            if self.model_config.hf_config.architectures == [
+                "MllamaForConditionalGeneration"
+            ]:
+                logger.info("Automatically turn off --chunked-prefill-size for mllama.")
+                server_args.chunked_prefill_size = -1
+
+            if self.model_config.hf_config.architectures == [
+                "Qwen2VLForConditionalGeneration"
+            ]:
+                # TODO: qwen2-vl does not support radix cache now, set disable_radix_cache=True automatically
+                logger.info(
+                    "Automatically turn off --chunked-prefill-size and disable radix cache for qwen2-vl."
+                )
+                server_args.chunked_prefill_size = -1
+                server_args.disable_radix_cache = True
+
     def init_torch_distributed(self):
         logger.info("Init torch distributed begin.")
-        torch.get_device_module(self.device).set_device(self.gpu_id)
 
+        torch.get_device_module(self.device).set_device(self.gpu_id)
         if self.device == "cuda":
             backend = "nccl"
         elif self.device == "xpu":
@@ -400,6 +395,18 @@ class ModelRunner:
             f"mem usage={(before_avail_memory - after_avail_memory):.2f} GB."
         )
 
+        # Handle the case where some ranks do not finish loading.
+        try:
+            dist.monitored_barrier(
+                group=get_tp_group().cpu_group,
+                timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S),
+                wait_all_ranks=True,
+            )
+        except RuntimeError:
+            raise ValueError(
+                f"TP rank {self.tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node."
+            ) from None
+
     def update_weights_from_disk(
         self, model_path: str, load_format: str
     ) -> tuple[bool, str]:
@@ -710,15 +717,6 @@ class ModelRunner:
             # Draft worker shares req_to_token_pool with the target worker.
             assert self.is_draft_worker
 
-        if self.token_to_kv_pool_allocator is None:
-            self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
-                self.max_total_num_tokens,
-                dtype=self.kv_cache_dtype,
-                device=self.device,
-            )
-        else:
-            assert self.is_draft_worker
-
         if (
             self.model_config.attention_arch == AttentionArch.MLA
             and not self.server_args.disable_mla
@@ -753,6 +751,17 @@ class ModelRunner:
                 device=self.device,
                 enable_memory_saver=self.server_args.enable_memory_saver,
             )
+
+        if self.token_to_kv_pool_allocator is None:
+            self.token_to_kv_pool_allocator = TokenToKVPoolAllocator(
+                self.max_total_num_tokens,
+                dtype=self.kv_cache_dtype,
+                device=self.device,
+                kvcache=self.token_to_kv_pool,
+            )
+        else:
+            assert self.is_draft_worker
+
         logger.info(
             f"Memory pool end. "
             f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
@@ -770,6 +779,10 @@ class ModelRunner:
     def init_attention_backend(self):
         """Init attention kernel backend."""
         if self.server_args.attention_backend == "flashinfer":
+            # Init streams
+            if self.server_args.speculative_algorithm == "EAGLE":
+                self.plan_stream_for_flashinfer = torch.cuda.Stream()
+
             self.attn_backend = FlashInferAttnBackend(self)
         elif self.server_args.attention_backend == "triton":
             assert self.sliding_window_size is None, (
@@ -878,18 +891,24 @@ class ModelRunner:
             forward_batch.input_ids, forward_batch.positions, forward_batch
         )
 
-    def forward(self, forward_batch: ForwardBatch) -> LogitsProcessorOutput:
+    def forward(
+        self, forward_batch: ForwardBatch, skip_attn_backend_init: bool = False
+    ) -> LogitsProcessorOutput:
         if (
             forward_batch.forward_mode.is_cuda_graph()
             and self.cuda_graph_runner
             and self.cuda_graph_runner.can_run(forward_batch)
         ):
-            return self.cuda_graph_runner.replay(forward_batch)
+            return self.cuda_graph_runner.replay(
+                forward_batch, skip_attn_backend_init=skip_attn_backend_init
+            )
 
         if forward_batch.forward_mode.is_decode():
             return self.forward_decode(forward_batch)
         elif forward_batch.forward_mode.is_extend():
-            return self.forward_extend(forward_batch)
+            return self.forward_extend(
+                forward_batch, skip_attn_backend_init=skip_attn_backend_init
+            )
         elif forward_batch.forward_mode.is_idle():
             return self.forward_idle(forward_batch)
         else:

sglang/srt/server_args.py

@@ -71,7 +71,6 @@ class ServerArgs:
     schedule_policy: str = "fcfs"
     schedule_conservativeness: float = 1.0
     cpu_offload_gb: int = 0
-    prefill_only_one_req: bool = False
 
     # Other runtime options
     tp_size: int = 1
@@ -277,19 +276,17 @@ class ServerArgs:
             self.speculative_algorithm = "EAGLE"
 
         if self.speculative_algorithm == "EAGLE":
-            self.disable_overlap_schedule = True
-            self.prefill_only_one_req = True
-            self.disable_cuda_graph_padding = True
             if self.max_running_requests is None:
                 self.max_running_requests = 32
+            self.disable_overlap_schedule = True
+            self.disable_cuda_graph_padding = True
             logger.info(
                 "Overlap scheduler are disabled because of using "
                 "eagle speculative decoding."
-                "Max running request set to 32 because of using eagle speculative decoding."
             )
             # The token generated from the verify step is counted.
             # If sepculative_num_steps >= speculative_num_draft_tokens, the additional tokens will definitely be discarded.
-            assert self.speculative_num_steps < self.speculative_num_draft_tokens
+            # assert self.speculative_num_steps < self.speculative_num_draft_tokens
 
         # GGUF
         if (
@@ -509,12 +506,6 @@ class ServerArgs:
             default=ServerArgs.cpu_offload_gb,
             help="How many GBs of RAM to reserve for CPU offloading",
         )
-        parser.add_argument(
-            "--prefill-only-one-req",
-            type=bool,
-            help="If true, we only prefill one request at one prefill batch",
-            default=ServerArgs.prefill_only_one_req,
-        )
 
         # Other runtime options
         parser.add_argument(

sglang/srt/speculative/build_eagle_tree.py

@@ -26,7 +26,12 @@ def build_tree_kernel_efficient_preprocess(
 
     draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
     draft_tokens = torch.cat((verified_id.unsqueeze(1), draft_tokens), dim=1).flatten()
-    parent_list = torch.cat(parents_list[:-1], 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
 

sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import bisect
-import time
 from typing import TYPE_CHECKING, Callable
 
 import torch
@@ -162,20 +161,11 @@ class EAGLEDraftCudaGraphRunner:
 
             run_once()
 
-            torch.cuda.synchronize()
-            self.model_runner.tp_group.barrier()
-
-        torch.cuda.synchronize()
-        self.model_runner.tp_group.barrier()
-
         with torch.cuda.graph(
             graph, pool=get_global_graph_memory_pool(), stream=stream
         ):
             out = run_once()
 
-        torch.cuda.synchronize()
-        self.model_runner.tp_group.barrier()
-
         set_global_graph_memory_pool(graph.pool())
         return graph, out
 
@@ -204,7 +194,7 @@ class EAGLEDraftCudaGraphRunner:
 
         # Attention backend
         self.model_runner.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
-            forward_batch
+            forward_batch, forward_batch.batch_size
         )
 
         # Replay

sglang/srt/speculative/eagle_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, List
 
 import torch
 import torch.nn.functional as F
@@ -62,6 +62,7 @@ class EagleDraftInput:
             batch.input_ids[pt : pt + extend_len] = torch.concat(
                 (input_ids[1:], self.verified_id[i].reshape(1))
             )
+            pt += extend_len
 
     def prepare_extend_after_decode(self, batch: ScheduleBatch, speculative_num_steps):
         assert self.verified_id.numel() == batch.out_cache_loc.shape[0]

sglang/srt/speculative/eagle_worker.py

@@ -1,20 +1,19 @@
 import logging
 import os
 import time
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple
 
 import torch
 from huggingface_hub import snapshot_download
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
     ForwardMode,
 )
-from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
     EAGLEDraftCudaGraphRunner,
@@ -27,7 +26,6 @@ from sglang.srt.speculative.eagle_utils import (
     fast_topk,
     select_top_k_tokens,
 )
-from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import get_available_gpu_memory
 
 logger = logging.getLogger(__name__)
@@ -44,16 +42,30 @@ class EAGLEWorker(TpModelWorker):
         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.padded_static_len = self.speculative_num_steps + 1
+        self.enable_nan_detection = server_args.enable_nan_detection
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self.target_worker = target_worker
+
         # Override context length with target model's context length
         server_args.context_length = target_worker.model_runner.model_config.context_len
-        os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"] = "1"
 
         # Do not capture cuda graph in `super().__init__()`
-        # We will capture it later
+        # It will be captured later.
         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()
+        )
 
-        # Lossy optimization by using hot tokens
+        # Load hot token ids
         if server_args.speculative_token_map is not None:
             self.hot_token_id = load_token_map(server_args.speculative_token_map)
             server_args.json_model_override_args = (
@@ -62,13 +74,7 @@ class EAGLEWorker(TpModelWorker):
         else:
             self.hot_token_id = None
 
-        # We share the allocator with a target worker. Draft/target worker
-        # owns its own KV cache.
-        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
-            target_worker.get_memory_pool()
-        )
-
-        # Init target worker
+        # Init draft worker
         super().__init__(
             gpu_id=gpu_id,
             tp_rank=tp_rank,
@@ -79,18 +85,6 @@ class EAGLEWorker(TpModelWorker):
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
         )
-        self.target_worker = target_worker
-
-        # Parse arguments
-        self.topk = server_args.speculative_eagle_topk
-        self.speculative_num_steps = server_args.speculative_num_steps
-        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
-            server_args.speculative_algorithm
-        )
-        self.server_args = server_args
-        self.use_nan_detection = self.server_args.enable_nan_detection
-        self.device = self.model_runner.device
-        self.gpu_id = self.model_runner.gpu_id
 
         # Share the embedding and lm_head
         embed, head = self.target_worker.model_runner.model.get_embed_and_head()
@@ -103,8 +97,12 @@ class EAGLEWorker(TpModelWorker):
             backup_disable_cuda_graph
         )
 
+        self.init_attention_backend()
+        self.init_cuda_graphs()
+
+    def init_attention_backend(self):
         # Create multi-step attn backends and cuda graph runners
-        if server_args.attention_backend == "flashinfer":
+        if self.server_args.attention_backend == "flashinfer":
             from sglang.srt.layers.attention.flashinfer_backend import (
                 FlashInferMultiStepDraftBackend,
             )
@@ -114,7 +112,7 @@ class EAGLEWorker(TpModelWorker):
                 self.topk,
                 self.speculative_num_steps,
             )
-        elif server_args.attention_backend == "triton":
+        elif self.server_args.attention_backend == "triton":
             from sglang.srt.layers.attention.triton_backend import (
                 TritonMultiStepDraftBackend,
             )
@@ -126,11 +124,9 @@ class EAGLEWorker(TpModelWorker):
             )
         else:
             raise ValueError(
-                f"EAGLE is not supportted in attention backend {server_args.attention_backend}"
+                f"EAGLE is not supportted in attention backend {self.server_args.attention_backend}"
             )
-
         self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
-        self.init_cuda_graphs()
 
     def init_cuda_graphs(self):
         """Capture cuda graphs."""
@@ -356,6 +352,41 @@ class EAGLEWorker(TpModelWorker):
         batch.forward_mode = ForwardMode.DECODE
         batch.spec_info = res.draft_input
 
+        if batch.return_logprob:
+            # Compute output logprobs using the sampler.
+            num_tokens_per_req = [
+                accept + 1 for accept in res.accept_length_per_req_cpu
+            ]
+            self.target_worker.model_runner.update_output_logprobs(
+                logits_output,
+                batch.sampling_info,
+                batch.top_logprobs_nums,
+                batch.token_ids_logprobs,
+                res.verified_id,
+                # +1 for bonus token.
+                num_tokens_per_req=num_tokens_per_req,
+            )
+
+            # Add output logprobs to the request.
+            pt = 0
+            # NOTE: tolist() of these values are skipped when output is processed
+            next_token_logprobs = res.logits_output.next_token_logprobs.tolist()
+            verified_ids = res.verified_id.tolist()
+            for req, num_tokens in zip(batch.reqs, num_tokens_per_req):
+                for _ in range(num_tokens):
+                    if req.return_logprob:
+                        token_id = verified_ids[pt]
+                        req.output_token_logprobs_val.append(next_token_logprobs[pt])
+                        req.output_token_logprobs_idx.append(token_id)
+                        if req.top_logprobs_num > 0:
+                            req.output_top_logprobs_val.append(
+                                res.logits_output.next_token_top_logprobs_val[pt]
+                            )
+                            req.output_top_logprobs_idx.append(
+                                res.logits_output.next_token_top_logprobs_idx[pt]
+                            )
+                    pt += 1
+
         return logits_output, res, model_worker_batch
 
     def forward_draft_extend(
@@ -381,6 +412,7 @@ class EAGLEWorker(TpModelWorker):
         forward_batch = ForwardBatch.init_new(
             model_worker_batch, self.draft_model_runner
         )
+        forward_batch.return_logprob = False
         logits_output = self.draft_model_runner.forward(forward_batch)
         self._detect_nan_if_needed(logits_output)
         assert isinstance(forward_batch.spec_info, EagleDraftInput)
@@ -393,6 +425,8 @@ class EAGLEWorker(TpModelWorker):
         batch.spec_info.prepare_extend_after_decode(batch, self.speculative_num_steps)
         batch.spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
         # We don't need logprob for this extend.
+        original_return_logprob = batch.return_logprob
+        batch.return_logprob = False
         model_worker_batch = batch.get_model_worker_batch()
         forward_batch = ForwardBatch.init_new(
             model_worker_batch, self.draft_model_runner
@@ -404,6 +438,7 @@ class EAGLEWorker(TpModelWorker):
 
         # Restore backup.
         # This is because `seq_lens` can be modified in `prepare_extend_after_decode`
+        batch.return_logprob = original_return_logprob
         batch.forward_mode = ForwardMode.DECODE
         batch.seq_lens = seq_lens_backup
 
@@ -415,7 +450,7 @@ class EAGLEWorker(TpModelWorker):
         draft_input.hidden_states = logits_output.hidden_states
 
     def _detect_nan_if_needed(self, logits_output: LogitsProcessorOutput):
-        if self.use_nan_detection:
+        if self.enable_nan_detection:
             logits = logits_output.next_token_logits
             if torch.any(torch.isnan(logits)):
                 logger.warning("Detected errors during sampling! NaN in the logits.")

sglang/version.py

@@ -1 +1 @@
-__version__ = "0.4.3.post3"
+__version__ = "0.4.3.post4"

{sglang-0.4.3.post3.dist-info → sglang-0.4.3.post4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sglang
-Version: 0.4.3.post3
+Version: 0.4.3.post4
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -239,6 +239,7 @@ Requires-Dist: xgrammar==0.1.14; extra == "runtime-common"
 Requires-Dist: ninja; extra == "runtime-common"
 Requires-Dist: transformers==4.48.3; extra == "runtime-common"
 Requires-Dist: llguidance>=0.6.15; extra == "runtime-common"
+Requires-Dist: datasets; extra == "runtime-common"
 Provides-Extra: srt
 Requires-Dist: sglang[runtime_common]; extra == "srt"
 Requires-Dist: sgl-kernel==0.0.3.post6; extra == "srt"